Update app.py

app.py

@@ -1,22 +1,40 @@
 from huggingface_hub import InferenceClient
-from langchain_community.llms import HuggingFaceHub
 from langchain_community.tools import DuckDuckGoSearchResults
 from langchain.agents import create_react_agent, AgentExecutor
 from langchain_core.tools import BaseTool
 from pydantic import Field
 from PIL import Image, ImageDraw, ImageFont
-import tempfile
+from functools import lru_cache
 import gradio as gr
 from io import BytesIO
-from typing import Optional
-from langchain_core.language_models.llms import LLM
 from transformers import pipeline
+from langchain_core.language_models.llms import LLM
+import os
+
+# === Global Model Setup ===
+
+# Preload image generation inference client
+image_client = InferenceClient("m-ric/text-to-image")
+
+# Preload text generation model via HuggingFace Transformers
+text_gen_pipeline = pipeline("text-generation", model="Qwen/Qwen2.5-72B-Instruct", max_new_tokens=512)
+
+# === LangChain Wrapper for Pipeline ===
+class PipelineLLM(LLM):
+    def _call(self, prompt, stop=None):
+        return text_gen_pipeline(prompt)[0]["generated_text"]
+
+    @property
+    def _llm_type(self):
+        return "pipeline_llm"
+
+llm = PipelineLLM()
 
-# === Image generation tool ===
+# === Image Tool ===
 class TextToImageTool(BaseTool):
     name: str = "text_to_image"
     description: str = "Generate an image from a text prompt."
-    client: InferenceClient = Field(exclude=True)
+    client: InferenceClient = Field(default=image_client, exclude=True)
 
     def _run(self, prompt: str) -> Image.Image:
         print(f"[Tool] Generating image for prompt: {prompt}")
@@ -26,16 +44,23 @@ class TextToImageTool(BaseTool):
     def _arun(self, prompt: str):
         raise NotImplementedError("This tool does not support async.")
 
+# Instantiate tools
+text_to_image_tool = TextToImageTool()
+search_tool = DuckDuckGoSearchResults()
+
+# Create LangChain agent
+agent = create_react_agent(llm=llm, tools=[text_to_image_tool, search_tool])
+agent_executor = AgentExecutor(agent=agent, tools=[text_to_image_tool, search_tool], verbose=True)
 
-# === Labeling Function ===
+# === Utility: Add Label to Image ===
 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(font_path, font_size)
+        font = ImageFont.truetype(font_path, 30)
     except:
         font = ImageFont.load_default()
+
     text_width, text_height = draw.textsize(label, font=font)
     position = (image.width - text_width - 20, image.height - text_height - 20)
     rect_position = [position[0] - 10, position[1] - 10, position[0] + text_width + 10, position[1] + text_height + 10]
@@ -43,8 +68,8 @@ def add_label_to_image(image, label):
     draw.text(position, label, fill="white", font=font)
     return image
 
-
 # === Prompt Generator ===
+@lru_cache(maxsize=128)
 def generate_prompts_for_object(object_name):
     return {
         "past": f"Show an old version of a {object_name} from its early days.",
@@ -52,54 +77,32 @@ def generate_prompts_for_object(object_name):
         "future": f"Show a futuristic version of a {object_name}, predicting future features/designs.",
     }
 
-
-# === Agent Setup ===
-# Set up the tools
-text_to_image_client = InferenceClient("m-ric/text-to-image")
-text_to_image_tool = TextToImageTool(client=text_to_image_client)
-search_tool = DuckDuckGoSearchResults()
-
-# Load a public, token-free model locally via transformers pipeline
-text_gen_pipeline = pipeline("text-generation", model="Qwen/Qwen2.5-72B-Instruct", max_new_tokens=512)
-#tiiuae/falcon-7b-instruct
-
-# Wrap pipeline into a LangChain LLM
-class PipelineLLM(LLM):
-    def _call(self, prompt, stop=None):
-        output = text_gen_pipeline(prompt)[0]["generated_text"]
-        return output
-
-    @property
-    def _llm_type(self):
-        return "pipeline_llm"
-
-llm = PipelineLLM()
-
-# Create agent and executor
-agent = create_react_agent(llm=llm, tools=[text_to_image_tool, search_tool])
-agent_executor = AgentExecutor(agent=agent, tools=[text_to_image_tool, search_tool], verbose=True)
-
-
 # === History Generator ===
+@lru_cache(maxsize=64)
+def generate_image_for_prompt(prompt, label):
+    img = text_to_image_tool._run(prompt)
+    return add_label_to_image(img, label)
+
 def generate_object_history(object_name: str):
     prompts = generate_prompts_for_object(object_name)
     images = []
-    labels = {
-        "past": f"{object_name} - Past",
-        "present": f"{object_name} - Present",
-        "future": f"{object_name} - Future"
-    }
+    file_paths = []
+
     for period, prompt in prompts.items():
-        result = text_to_image_tool._run(prompt)
-        labeled = add_label_to_image(result, labels[period])
-        file_path = f"{object_name}_{period}.png"
-        labeled.save(file_path)
-        images.append((file_path, labels[period]))
-    gif_path = f"{object_name}_evolution.gif"
-    pil_images = [Image.open(img[0]) for img in images]
+        label = f"{object_name} - {period.capitalize()}"
+        labeled_image = generate_image_for_prompt(prompt, label)
+
+        file_path = f"/tmp/{object_name}_{period}.png"
+        labeled_image.save(file_path)
+        images.append((file_path, label))
+        file_paths.append(file_path)
+
+    # Create GIF
+    gif_path = f"/tmp/{object_name}_evolution.gif"
+    pil_images = [Image.open(p) for p in file_paths]
     pil_images[0].save(gif_path, save_all=True, append_images=pil_images[1:], duration=1000, loop=0)
-    return images, gif_path
 
+    return images, gif_path
 
 # === Gradio UI ===
 def create_gradio_interface():
@@ -117,7 +120,7 @@ def create_gradio_interface():
 
     return demo
 
-
 # === Launch App ===
-demo = create_gradio_interface()
-demo.launch(share=True)
+if __name__ == "__main__":
+    demo = create_gradio_interface()
+    demo.launch(share=True)