Update app.py

app.py

@@ -1,30 +1,14 @@
 from PIL import Image, ImageDraw, ImageFont
+import tempfile
 import gradio as gr
-from smolagents import CodeAgent, InferenceClientModel, DuckDuckGoSearchTool, Tool
-from gradio_client import Client
+from smolagents import CodeAgent, InferenceClientModel
+from smolagents import DuckDuckGoSearchTool, Tool
+from huggingface_hub import InferenceClient
 
-#%% Tool Wrapper for the Hugging Face Space
-class TextToImageTool(Tool):
-    name = "text_to_image"
-    description = "Generate an image from a text prompt using m-ric/text-to-image."
-    
-    inputs = {
-        "prompt": {
-            "type": "string",
-            "description": "Text prompt to generate the image"
-        }
-    }
-
-    output_type = "image"
-
-    def __init__(self):
-        super().__init__()
-        self.client = Client("m-ric/text-to-image")
-
-    def run(self, prompt):  # Must explicitly match 'inputs' keys
-        return self.client.predict(prompt, api_name="/predict")
+# =========================================================
+# Utility functions
+# =========================================================
 
-#%% Utility functions
 def add_label_to_image(image, label):
     draw = ImageDraw.Draw(image)
     font_path = "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"
@@ -33,25 +17,33 @@ def add_label_to_image(image, label):
         font = ImageFont.truetype(font_path, font_size)
     except:
         font = ImageFont.load_default()
+
     text_bbox = draw.textbbox((0, 0), label, font=font)
-    text_width = text_bbox[2] - text_bbox[0]
-    text_height = text_bbox[3] - text_bbox[1]
+    text_width, text_height = text_bbox[2] - text_bbox[0], text_bbox[3] - text_bbox[1]
     position = (image.width - text_width - 20, image.height - text_height - 20)
+
     rect_margin = 10
     rect_position = [
-        position[0] - rect_margin, position[1] - rect_margin,
-        position[0] + text_width + rect_margin, position[1] + text_height + rect_margin
+        position[0] - rect_margin,
+        position[1] - rect_margin,
+        position[0] + text_width + rect_margin,
+        position[1] + text_height + rect_margin,
     ]
     draw.rectangle(rect_position, fill=(0, 0, 0, 128))
     draw.text(position, label, fill="white", font=font)
     return image
 
-def plot_and_save_agent_image(image, label, save_path=None):
-    labeled_image = add_label_to_image(image, label)
+
+def plot_and_save_agent_image(agent_image, label, save_path=None):
+    pil_image = agent_image.to_raw()
+    labeled_image = add_label_to_image(pil_image, label)
     labeled_image.show()
     if save_path:
         labeled_image.save(save_path)
         print(f"Image saved to {save_path}")
+    else:
+        print("No save path provided. Image not saved.")
+
 
 def generate_prompts_for_object(object_name):
     return {
@@ -60,7 +52,43 @@ def generate_prompts_for_object(object_name):
         "future": f"Show a futuristic version of a {object_name}, by predicting advanced features and futuristic design."
     }
 
+# =========================================================
+# Tool wrapper for m-ric/text-to-image
+# =========================================================
+
+class WrappedTextToImageTool(Tool):
+    name = "text_to_image"
+    description = "Generates an image from a text prompt using the m-ric/text-to-image tool."
+    inputs = {
+        "prompt": {
+            "type": "string",
+            "description": "Text prompt to generate an image"
+        }
+    }
+    output_type = "image"
+
+    def __init__(self):
+        self.client = InferenceClient("m-ric/text-to-image")
+
+    def forward(self, prompt):
+        return self.client.text_to_image(prompt)
+
+# =========================================================
+# Tool and Agent Initialization
+# =========================================================
+
+image_generation_tool = WrappedTextToImageTool()
+search_tool = DuckDuckGoSearchTool()
+llm_engine = InferenceClientModel("Qwen/Qwen2.5-72B-Instruct")
+
+agent = CodeAgent(tools=[image_generation_tool, search_tool], model=llm_engine)
+
+# =========================================================
+# Main logic for image generation
+# =========================================================
+
 def generate_object_history(object_name):
+    images = []
     prompts = generate_prompts_for_object(object_name)
     labels = {
         "past": f"{object_name} - Past",
@@ -68,33 +96,29 @@ def generate_object_history(object_name):
         "future": f"{object_name} - Future"
     }
 
-    images = []
     for time_period, prompt in prompts.items():
         print(f"Generating {time_period} frame: {prompt}")
-        result = agent.run(prompt)  # Runs tool
-        if hasattr(result, "to_raw"):  # If wrapped output
-            result = result.to_raw()
-        images.append(result)
-        plot_and_save_agent_image(result, labels[time_period], save_path=f"{object_name}_{time_period}.png")
+        result = agent.run(prompt)
+        images.append(result.to_raw())
+        image_filename = f"{object_name}_{time_period}.png"
+        plot_and_save_agent_image(result, labels[time_period], save_path=image_filename)
 
     gif_path = f"{object_name}_evolution.gif"
     images[0].save(gif_path, save_all=True, append_images=images[1:], duration=1000, loop=0)
-    return images, gif_path
-
-#%% Tool & Agent Setup
-image_generation_tool = TextToImageTool()
-search_tool = DuckDuckGoSearchTool()
-llm_engine = InferenceClientModel("Qwen/Qwen2.5-72B-Instruct")
+    return [(f"{object_name}_past.png", labels["past"]),
+            (f"{object_name}_present.png", labels["present"]),
+            (f"{object_name}_future.png", labels["future"])], gif_path
 
-agent = CodeAgent(tools=[image_generation_tool, search_tool], model=llm_engine)
+# =========================================================
+# Gradio Interface
+# =========================================================
 
-#%% Gradio Interface
 def create_gradio_interface():
     with gr.Blocks() as demo:
-        gr.Markdown("# TimeMetamorphy: an object Evolution Generator")
+        gr.Markdown("# TimeMetamorphy: An Object Evolution Generator")
         gr.Markdown("""
-        ## Unlocking the secrets of time!
-        Enter an object name (like bicycle or smartphone), and this app will generate its visual evolution.
+        Explore how everyday objects evolved over time. Enter an object name like "phone", "car", or "bicycle"
+        and see its past, present, and future visualized with AI!
         """)
 
         default_images = [
@@ -106,14 +130,18 @@ def create_gradio_interface():
 
         with gr.Row():
             with gr.Column():
-                object_name_input = gr.Textbox(label="Enter an object name", placeholder="e.g., bicycle, phone")
+                object_name_input = gr.Textbox(label="Enter an object name", placeholder="e.g. bicycle, car, phone")
                 generate_button = gr.Button("Generate Evolution")
                 image_gallery = gr.Gallery(label="Generated Images", columns=3, rows=1, value=default_images)
                 gif_output = gr.Image(label="Generated GIF", value=default_gif_path)
 
         generate_button.click(fn=generate_object_history, inputs=[object_name_input], outputs=[image_gallery, gif_output])
+
     return demo
 
-# Launch app
+# =========================================================
+# Run the app
+# =========================================================
+
 demo = create_gradio_interface()
 demo.launch(share=True)