Spaces:
Runtime error
Runtime error
Update app.py
Browse files
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
|
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 |
|