Update app.py

app.py

@@ -10,19 +10,21 @@ from tools.final_answer import FinalAnswerTool
 from Gradio_UI import GradioUI
 
 # Below is an example of a tool that does nothing. Amaze us with your creativity !
+from smolagents.agent_types import AgentImage  # 🔑 Импортируем AgentImage
+
 @tool
 def generate_image(
     prompt: str, 
     negative_prompt: str = "blurry, text, watermark, low quality", 
     width: int = 1024, 
     height: int = 1024
-) -> Image.Image:
+) -> AgentImage:  # 🔑 Указываем тип возвращаемого значения AgentImage
     """
     Генерирует изображение по текстовому описанию.
     
     Пример:
         >>> generate_image("cute kitten on a blanket")
-        <PIL.Image.Image>
+        <AgentImage>
     
     Args:
         prompt: Что нарисовать (подробно, на английском лучше).
@@ -31,14 +33,13 @@ def generate_image(
         height: Высота в пикселях (кратно 8, макс 1024).
     
     Returns:
-        PIL.Image.Image: Сгенерированное изображение.
+        AgentImage: Сгенерированное изображение.
     """
     from huggingface_hub import InferenceClient
     from PIL import Image
     import io
-    import os, sys
+    import os, tempfile, sys
     
-    # 🔍 Отладка
     print(f"🔍 generate_image: prompt='{prompt[:50]}...', size={width}x{height}", file=sys.stderr)
     
     client = InferenceClient(token=os.getenv("HF_TOKEN"))
@@ -54,22 +55,21 @@ def generate_image(
             guidance_scale=3.5
         )
         
-        # 🔑 Конвертация в PIL.Image
+        # Получаем PIL объект
         if isinstance(result, Image.Image):
             image = result
-        elif isinstance(result, bytes):
-            image = Image.open(io.BytesIO(result))
-        elif hasattr(result, 'image') and isinstance(result.image, Image.Image):
-            image = result.image  # Распаковка AgentImage
         else:
-            # Попытка обработать как поток байтов
             image = Image.open(io.BytesIO(result))
         
-        print(f"✅ generate_image: изображение успешно сгенерировано!", file=sys.stderr)
+        # 🔑 Возвращаем сохранение во временный файл
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
+            image.save(tmp.name, format="PNG")
+            file_path = tmp.name
+            
+        print(f"✅ generate_image: сохранено в {file_path}", file=sys.stderr)
         
-        # 🔑 ВОЗВРАЩАЕМ САМ ОБЪЕКТ PIL.Image!
-        # GradioUI сам поймет, что это картинка, и правильно отобразит её в чате
-        return image
+        # 🔑 Возвращаем AgentImage, передав ему путь к сохраненному файлу!
+        return AgentImage(file_path)
         
     except Exception as e:
         print(f"❌ generate_image exception: {e}", file=sys.stderr)