Update app.py

app.py

@@ -3,6 +3,7 @@ import datetime
 import requests
 import pytz
 import yaml
+import torch
 from tools.final_answer import FinalAnswerTool
 
 from Gradio_UI import GradioUI
@@ -33,6 +34,44 @@ def get_current_time_in_timezone(timezone: str) -> str:
     except Exception as e:
         return f"Error fetching time for timezone '{timezone}': {str(e)}"
 
+@tool
+def generate_open_source_image(prompt: str) -> str:
+    """
+    Genera un'immagine a partire da un prompt di testo utilizzando un modello open source di diffusion (Stable Diffusion).
+    
+    Args:
+        prompt: Una stringa contenente il prompt testuale da usare per generare l'immagine.
+        
+    Returns:
+        Una stringa che rappresenta il percorso del file in cui l'immagine generata è stata salvata.
+    """
+    try:
+        # Importa il pipeline di Stable Diffusion
+        from diffusers import StableDiffusionPipeline
+
+        # Specifica l'ID del modello open source su Hugging Face
+        model_id = "stabilityai/stable-diffusion-2-1"
+        
+        # Determina il dispositivo da usare (GPU se disponibile, altrimenti CPU)
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        
+        # Carica il modello; utilizza torch.float16 se disponibile la GPU per prestazioni migliori
+        pipe = StableDiffusionPipeline.from_pretrained(
+            model_id, 
+            torch_dtype=torch.float16 if device == "cuda" else torch.float32
+        )
+        pipe = pipe.to(device)
+        
+        # Genera l'immagine a partire dal prompt
+        image = pipe(prompt).images[0]
+        
+        # Crea un nome di file unico per salvare l'immagine
+        filename = f"generated_image_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.png"
+        image.save(filename)
+        
+        return f"Immagine generata e salvata in: {filename}"
+    except Exception as e:
+        return f"Si è verificato un errore nella generazione dell'immagine: {str(e)}"
 
 final_answer = FinalAnswerTool()
 
@@ -55,7 +94,8 @@ with open("prompts.yaml", 'r') as stream:
     
 agent = CodeAgent(
     model=model,
-    tools=[final_answer], ## add your tools here (don't remove final answer)
+    tools=[final_answer, 
+          generate_open_source_image], ## add your tools here (don't remove final answer)
     max_steps=6,
     verbosity_level=1,
     grammar=None,