Spaces:
Running
Running
| import gradio as gr | |
| from smolagents import CodeAgent, HfApiModel, Tool | |
| from PIL import Image | |
| import os | |
| # Define the image generation tool | |
| image_generation_tool = Tool.from_space( | |
| "black-forest-labs/FLUX.1-schnell", | |
| name="image_generator", | |
| description="Generate an image from a text prompt." | |
| ) | |
| # Initialize the model | |
| model = HfApiModel("Qwen/Qwen2.5-Coder-32B-Instruct") # Replace with your model | |
| # Create the Smol Agent | |
| agent = CodeAgent(tools=[image_generation_tool], model=model) | |
| # Define the chatbot function | |
| def chat_with_agent(user_input, history): | |
| if history is None: | |
| history = [] | |
| # Run the agent and debug the output | |
| response = agent.run(user_input) | |
| print(f"Agent response: {response} (type: {type(response)})") # Debug print | |
| # Initialize outputs | |
| image_output = None | |
| response_text = "" | |
| # Handle different response types | |
| if isinstance(response, str): | |
| # Check if it’s a file path | |
| if any(response.endswith(ext) for ext in [".webp", ".png", ".jpg", ".jpeg"]) and os.path.exists(response): | |
| try: | |
| image_output = Image.open(response) | |
| response_text = "Image generated! See below." | |
| except Exception as e: | |
| response_text = f"Error loading image: {str(e)}" | |
| else: | |
| # Treat as plain text | |
| response_text = response | |
| elif isinstance(response, int): # Handle unexpected integer response | |
| response_text = f"Agent returned a number ({response}), which I can’t display directly. Try asking differently!" | |
| else: | |
| # Fallback for unexpected types | |
| response_text = f"Unexpected response: {str(response)}. I’ll try to handle it better next time!" | |
| # Append to history as a valid [user, assistant] pair | |
| history.append([user_input, response_text]) | |
| return history, "", image_output | |
| # Create the Gradio interface | |
| with gr.Blocks(title="Smol Agent Chatbot") as demo: | |
| gr.Markdown("# Chat with a Smol Agent") | |
| gr.Markdown("Ask anything—text or image generation supported!") | |
| # Chatbot for text conversation | |
| chatbot = gr.Chatbot() | |
| # Image output for displaying generated images | |
| image_display = gr.Image(label="Generated Image") | |
| # Textbox for user input | |
| user_input = gr.Textbox(placeholder="Type your message here...", label="Your Message") | |
| # Submit button | |
| submit_btn = gr.Button("Send") | |
| # Connect components | |
| submit_btn.click( | |
| fn=chat_with_agent, | |
| inputs=[user_input, chatbot], | |
| outputs=[chatbot, user_input, image_display] | |
| ) | |
| # Launch the app | |
| demo.launch(debug=True) # Enable debug mode to see logs |