Update app.py

app.py

@@ -7,52 +7,54 @@ from PIL import Image, ImageDraw, ImageFont
 from functools import lru_cache
 import gradio as gr
 from io import BytesIO
-from transformers import pipeline
-from langchain_core.language_models.llms import LLM
 import os
 
-# === Global Model Setup ===
+# === Setup Inference Clients ===
+# Use your Hugging Face token if necessary:
+# client = InferenceClient(repo_id="model", token="YOUR_HF_TOKEN")
 
-# Preload image generation inference client
 image_client = InferenceClient("m-ric/text-to-image")
+text_client = InferenceClient("Qwen/Qwen2.5-72B-Instruct")
 
-# 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 using InferenceClient for text generation ===
+class InferenceClientLLM(BaseTool):
+    name: str = "inference_text_generator"
+    description: str = "Generate text using HF Inference API."
+    client: InferenceClient = Field(default=text_client, exclude=True)
 
-# === LangChain Wrapper for Pipeline ===
-class PipelineLLM(LLM):
-    def _call(self, prompt, stop=None):
-        return text_gen_pipeline(prompt)[0]["generated_text"]
+    def _run(self, prompt: str) -> str:
+        print(f"[LLM] Generating text for prompt: {prompt}")
+        response = self.client.text_generation(prompt)
+        # response is usually a dict with 'generated_text'
+        return response.get("generated_text", "")
 
-    @property
-    def _llm_type(self):
-        return "pipeline_llm"
-
-llm = PipelineLLM()
+    def _arun(self, prompt: str):
+        raise NotImplementedError("Async not supported.")
 
-# === Image Tool ===
+# === Image generation tool ===
 class TextToImageTool(BaseTool):
     name: str = "text_to_image"
     description: str = "Generate an image from a text prompt."
     client: InferenceClient = Field(default=image_client, exclude=True)
 
     def _run(self, prompt: str) -> Image.Image:
-        print(f"[Tool] Generating image for prompt: {prompt}")
+        print(f"[Image Tool] Generating image for prompt: {prompt}")
         image_bytes = self.client.text_to_image(prompt)
         return Image.open(BytesIO(image_bytes))
 
     def _arun(self, prompt: str):
-        raise NotImplementedError("This tool does not support async.")
+        raise NotImplementedError("Async not supported.")
 
-# Instantiate tools
+# === Initialize tools ===
 text_to_image_tool = TextToImageTool()
+text_gen_tool = InferenceClientLLM()
 search_tool = DuckDuckGoSearchResults()
 
-# Create LangChain agent
-agent = create_react_agent(llm=llm, tools=[text_to_image_tool, search_tool])
+# === Create agent ===
+agent = create_react_agent(llm=text_gen_tool, tools=[text_to_image_tool, search_tool])
 agent_executor = AgentExecutor(agent=agent, tools=[text_to_image_tool, search_tool], verbose=True)
 
-# === Utility: Add Label to Image ===
+# === Image labeling ===
 def add_label_to_image(image, label):
     draw = ImageDraw.Draw(image)
     font_path = "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"
@@ -68,7 +70,7 @@ def add_label_to_image(image, label):
     draw.text(position, label, fill="white", font=font)
     return image
 
-# === Prompt Generator ===
+# === Prompt generation with caching ===
 @lru_cache(maxsize=128)
 def generate_prompts_for_object(object_name):
     return {
@@ -77,12 +79,13 @@ def generate_prompts_for_object(object_name):
         "future": f"Show a futuristic version of a {object_name}, predicting future features/designs.",
     }
 
-# === History Generator ===
+# === Cache generated images ===
 @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)
 
+# === Main generation function ===
 def generate_object_history(object_name: str):
     prompts = generate_prompts_for_object(object_name)
     images = []
@@ -120,7 +123,7 @@ def create_gradio_interface():
 
     return demo
 
-# === Launch App ===
+# === Launch app ===
 if __name__ == "__main__":
     demo = create_gradio_interface()
     demo.launch(share=True)