imagine.py 1.3 KB

1234567891011121314151617181920212223242526272829303132333435363738
  1. import asyncio
  2. import functools
  3. from nio.events.room_events import RoomMessageText
  4. from nio.rooms import MatrixRoom
  5. async def command_imagine(room: MatrixRoom, event: RoomMessageText, bot):
  6. prompt = " ".join(event.body.split()[2:])
  7. if prompt:
  8. bot.logger.log("Generating image...")
  9. try:
  10. loop = asyncio.get_event_loop()
  11. except Exception as e:
  12. bot.logger.log(f"Error getting event loop: {e}", "error")
  13. await bot.send_message(
  14. room, "Something went wrong. Please try again.", True)
  15. return
  16. try:
  17. image_partial = functools.partial(bot.image_api.generate_image, prompt, user=room.room_id)
  18. images, tokens_used = await loop.run_in_executor(None, image_partial)
  19. except Exception as e:
  20. bot.logger.log(f"Error generating image: {e}", "error")
  21. await bot.send_message(room, "Sorry, I couldn't generate an image. Please try again later.", True)
  22. return
  23. for image in images:
  24. bot.logger.log(f"Sending image...")
  25. await bot.send_image(room, image)
  26. bot.log_api_usage(event, room, f"{bot.image_api.api_code}-{bot.image_api.image_api}", tokens_used)
  27. return
  28. await bot.send_message(room, "You need to provide a prompt.", True)