Update app.py

app.py

@@ -1,64 +1,119 @@
-import torch
+from transformers import HfAgent
 from PIL import Image, ImageDraw, ImageFont
 import gradio as gr
-from diffusers import StableDiffusionPipeline
+import os
 
-# Load Stable Diffusion model
-pipe = StableDiffusionPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16
-).to("cuda" if torch.cuda.is_available() else "cpu")
+#%% Utility Functions
 
-# Function to add label
 def add_label_to_image(image, label):
     draw = ImageDraw.Draw(image)
+    font_path = "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"
+    font_size = 30
     try:
-        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 30)
+        font = ImageFont.truetype(font_path, font_size)
     except:
         font = ImageFont.load_default()
-    position = (20, image.height - 50)
-    draw.rectangle([position, (position[0]+400, position[1]+40)], fill=(0, 0, 0, 180))
-    draw.text(position, label, font=font, fill="white")
+    
+    text_bbox = draw.textbbox((0, 0), label, font=font)
+    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,
+    ]
+    draw.rectangle(rect_position, fill=(0, 0, 0, 128))
+    draw.text(position, label, fill="white", font=font)
     return image
 
-# Generate prompt images
-def generate_object_history(object_name):
-    prompts = {
-        "past": f"An old version of a {object_name}, vintage, old-fashioned",
-        "present": f"A modern {object_name}, realistic, current design",
-        "future": f"A futuristic {object_name}, sci-fi, advanced design"
+def plot_and_save_agent_image(image, label, save_path=None):
+    labeled_image = add_label_to_image(image, label)
+    labeled_image.show()
+    if save_path:
+        labeled_image.save(save_path)
+        print(f"Image saved to {save_path}")
+
+def generate_prompts_for_object(object_name):
+    return {
+        "past": f"Show an old version of a {object_name} from its early days.",
+        "present": f"Show a {object_name} with current features/design/technology.",
+        "future": f"Show a futuristic version of a {object_name}, predicting advanced features and design."
     }
 
+#%% HF Agent Initialization
+
+agent = HfAgent("https://api-inference.huggingface.co/models/HuggingFaceH4/zephyr-7b-beta")
+
+#%% Core Generation Function
+
+def generate_object_history(object_name):
     images = []
-    pil_images = []
+    gif_frames = []
+
+    prompts = generate_prompts_for_object(object_name)
+    labels = {
+        "past": f"{object_name} - Past",
+        "present": f"{object_name} - Present",
+        "future": f"{object_name} - Future"
+    }
+
+    output_image_paths = []
 
     for period, prompt in prompts.items():
-        image = pipe(prompt).images[0]
-        labeled_image = add_label_to_image(image, f"{object_name.title()} - {period.title()}")
+        print(f"Generating image for: {prompt}")
+        result = agent.run(prompt)
+        
+        # Find first image in result (agent returns a dict)
+        image_path = next((v for v in result.values() if isinstance(v, str) and v.endswith((".png", ".jpg"))), None)
+        if not image_path or not os.path.exists(image_path):
+            raise RuntimeError(f"No valid image generated for {prompt}")
+
+        image = Image.open(image_path).convert("RGB")
+        labeled_image = add_label_to_image(image, labels[period])
+
         filename = f"{object_name}_{period}.png"
         labeled_image.save(filename)
-        images.append((filename, f"{object_name.title()} - {period.title()}"))
-        pil_images.append(labeled_image)
+        output_image_paths.append((filename, labels[period]))
+        gif_frames.append(labeled_image)
 
+    # Save animated GIF
     gif_path = f"{object_name}_evolution.gif"
-    pil_images[0].save(gif_path, save_all=True, append_images=pil_images[1:], duration=1000, loop=0)
+    gif_frames[0].save(gif_path, save_all=True, append_images=gif_frames[1:], duration=1000, loop=0)
 
-    return images, gif_path
+    return output_image_paths, gif_path
+
+#%% Gradio Interface
 
-# Gradio Interface
 def create_gradio_interface():
     with gr.Blocks() as demo:
-        gr.Markdown("# TimeMetamorphy: Object Evolution Visualizer")
+        gr.Markdown("# TimeMetamorphy: An Object Evolution Generator")
+
+        gr.Markdown("""
+        Explore how objects change over time — from past, to present, to future.
+        Enter any object below and let AI visualize its transformation through the ages.
+        """)
+
+        default_images = [
+            ("car_past.png", "Car - Past"),
+            ("car_present.png", "Car - Present"),
+            ("car_future.png", "Car - Future")
+        ]
+        default_gif_path = "car_evolution.gif"
 
-        object_name_input = gr.Textbox(label="Enter an object name (e.g., bicycle, phone)")
-        generate_button = gr.Button("Generate Evolution")
+        with gr.Row():
+            with gr.Column():
+                object_name_input = gr.Textbox(label="Enter an object name", placeholder="e.g., bicycle, phone")
+                generate_button = gr.Button("Generate Evolution")
+                image_gallery = gr.Gallery(label="Generated Images", show_label=True, columns=3, rows=1, value=default_images)
+                gif_output = gr.Image(label="Generated GIF", show_label=True, value=default_gif_path)
 
-        image_gallery = gr.Gallery(label="Generated Images", columns=3, rows=1)
-        gif_output = gr.Image(label="Generated GIF")
+        generate_button.click(fn=generate_object_history, inputs=[object_name_input], outputs=[image_gallery, gif_output])
 
-        generate_button.click(fn=generate_object_history,
-                              inputs=[object_name_input],
-                              outputs=[image_gallery, gif_output])
     return demo
 
+# Launch
 demo = create_gradio_interface()
 demo.launch(share=True)