Update app.py

app.py

@@ -4,7 +4,7 @@ import gradio as gr
 from PIL import Image
 import tempfile
 import shutil
-from functools import partial # To create a callable for our custom tool
+from functools import partial 
 
 from diffusers import StableDiffusionPipeline
 from huggingface_hub import InferenceClient
@@ -22,13 +22,12 @@ HF_TOKEN = os.environ.get("HF_TOKEN") # Using HF_TOKEN for consistency with Hugg
 
 # Define the model ID for image generation
 IMAGE_GEN_MODEL_ID = "segmind/tiny-sd" # Using the smaller model as it loaded successfully
-# IMAGE_GEN_MODEL_ID = "runwayml/stable-diffusion-v1-5" # You can try this again after proving basic functionality
 
 print(f"Loading Stable Diffusion Pipeline directly on GPU: {IMAGE_GEN_MODEL_ID}...")
 try:
     pipe = StableDiffusionPipeline.from_pretrained(
         IMAGE_GEN_MODEL_ID,
-        torch_dtype=torch.float16,
+        torch_dtype=torch.float16, # Use float16 for less VRAM usage on T4
         use_safetensors=False, # Set to False for models that don't have safetensors (like tiny-sd)
         token=HF_TOKEN # Pass token for potential faster model download
     )
@@ -80,8 +79,6 @@ if pipe is None:
     raise RuntimeError("Cannot start agent as image generation pipeline failed to load. Check logs.")
 
 # Instantiate the LLM for the agent
-# Using HuggingFaceHub to connect to Zephyr-7b-beta model on HF Inference API
-# Ensure HF_TOKEN is set as a Space Secret
 llm = HuggingFaceHub(
     repo_id="HuggingFaceH4/zephyr-7b-beta",
     huggingfacehub_api_token=HF_TOKEN, # Use HF_TOKEN directly as required by HuggingFaceHub LLM
@@ -97,6 +94,7 @@ prompt_template = ChatPromptTemplate.from_messages(
     [
         ("system", """You are a powerful AI assistant that can generate images and search the web.
 You have access to the following tools: {tools}
+Available tools: {tool_names} # <--- ADDED THIS LINE: Provide tool names to the LLM for ReAct agent.
 
 When you need to generate an image, use the `image_generator` tool. Its input must be a very detailed, descriptive text string.
 When you need factual information or context, use the `search` tool.
@@ -131,10 +129,6 @@ def run_agent_in_gradio(message, history):
         chat_history.append(AIMessage(content=ai_msg))
 
     try:
-        # Stream output from the agent
-        # LangChain AgentExecutor doesn't directly stream token by token in a simple loop
-        # For streaming, you'd typically use .stream() or a custom callback handler.
-        # For simplicity in Gradio ChatInterface, we'll run it once.
         response = agent_executor.invoke({"input": message, "chat_history": chat_history})
         agent_output = response["output"]