Update app.py

app.py

@@ -3,6 +3,8 @@ import datetime
 import requests
 import pytz
 import yaml
+import torch
+from diffusers import FluxPipeline
 from tools.final_answer import FinalAnswerTool
 
 from Gradio_UI import GradioUI
@@ -34,6 +36,20 @@ def get_current_time_in_timezone(timezone: str) -> str:
     except Exception as e:
         return f"Error fetching time for timezone '{timezone}': {str(e)}"
 
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
+pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
+
+prompt = "{local_time}"
+image = pipe(
+    prompt,
+    height=1024,
+    width=1024,
+    guidance_scale=3.5,
+    num_inference_steps=50,
+    max_sequence_length=512,
+    generator=torch.Generator("cpu").manual_seed(0)
+).images[0]
+image.save("flux-dev.png")
 
 final_answer = FinalAnswerTool()