Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import requests | |
| import base64 | |
| from PIL import Image | |
| import io | |
| import time | |
| import os | |
| # Your endpoint URL and token (you might want to use environment variables for these in production) | |
| API_URL = "https://ca80xvp8jmhqanbz.us-east-1.aws.endpoints.huggingface.cloud" | |
| def process_image(input_image, prompt): | |
| # Convert Gradio's image input to base64 | |
| if input_image is None: | |
| return None, "Please upload an image" | |
| # Convert PIL Image directly to base64 string | |
| buffer = io.BytesIO() | |
| input_image.save(buffer, format="WEBP") | |
| base64_image = f"data:image/webp;base64,{base64.b64encode(buffer.getvalue()).decode('utf-8')}" | |
| # Prepare headers and data | |
| headers = { | |
| "Authorization": f"Bearer {os.getenv('HF_API_KEY')}", | |
| "Content-Type": "application/json", | |
| "Accept": "application/json" | |
| } | |
| data = { | |
| "inputs": { | |
| "prompt": prompt, | |
| "image": base64_image | |
| } | |
| } | |
| # Make the API request with retry mechanism | |
| max_retries = 5 | |
| retry_delay = 60 # seconds | |
| for retry in range(max_retries): | |
| try: | |
| response = requests.post(API_URL, headers=headers, json=data, timeout=300) | |
| if response.status_code == 200: | |
| # Convert the response base64 directly to PIL Image | |
| result = response.json() | |
| image_bytes = base64.b64decode(result["final_image"]) | |
| output_image = Image.open(io.BytesIO(image_bytes)) | |
| # Return both the image and the parameters used | |
| params_text = "\nParameters used:\n" | |
| for key, value in result["parameters"].items(): | |
| params_text += f"{key}: {value}\n" | |
| return output_image, params_text | |
| elif response.status_code == 503: | |
| if retry < max_retries - 1: | |
| time.sleep(retry_delay) | |
| continue | |
| return None, f"Service unavailable after {max_retries} retries" | |
| else: | |
| return None, f"Error: {response.status_code}\n{response.text}" | |
| except requests.exceptions.RequestException as e: | |
| if retry < max_retries - 1: | |
| time.sleep(retry_delay) | |
| continue | |
| return None, f"Request error: {str(e)}" | |
| return None, "Maximum retries reached" | |
| # Create Gradio interface | |
| with gr.Blocks(title="Room Design Diffusion", theme=gr.themes.Soft()) as demo: | |
| gr.Markdown(""" | |
| # Room Design Diffusion | |
| Upload a room image and provide a prompt describing how you'd like to redesign it. | |
| The AI will generate a new version of your room based on your description. | |
| """) | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_image = gr.Image(label="Upload Room Image", type="pil") | |
| prompt = gr.Textbox( | |
| label="Describe how you want to redesign the room", | |
| placeholder="Example: A modern living room with a comfortable gray sectional sofa, glass coffee table, minimalist TV stand, and geometric area rug", | |
| lines=3 | |
| ) | |
| submit_btn = gr.Button("Generate Design", variant="primary") | |
| with gr.Column(): | |
| output_image = gr.Image(label="Generated Design") | |
| output_text = gr.Textbox(label="Processing Details", lines=5) | |
| submit_btn.click( | |
| fn=process_image, | |
| inputs=[input_image, prompt], | |
| outputs=[output_image, output_text] | |
| ) | |
| gr.Markdown(""" | |
| ### Tips for best results: | |
| - Use clear, detailed prompts describing the desired room style | |
| - Make sure the input image is well-lit and clearly shows the room | |
| - Be patient as generation can take a few minutes | |
| """) | |
| # Launch the app with sharing enabled | |
| demo.launch(share=True) |