bielas194 commited on
Commit
2d54ce1
Β·
verified Β·
1 Parent(s): 63621b6

Update app.py

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