Spaces:
Running
Running
| import gradio as gr | |
| import numpy as np | |
| import random | |
| import os | |
| import tempfile | |
| from PIL import Image, ImageOps | |
| import pillow_heif # For HEIF/AVIF support | |
| import io | |
| # --- Constants --- | |
| MAX_SEED = np.iinfo(np.int32).max | |
| def load_client(): | |
| """Initialize the Inference Client""" | |
| # Register HEIF opener with PIL for AVIF/HEIF support | |
| pillow_heif.register_heif_opener() | |
| # Get token from environment variable | |
| hf_token = os.getenv("HF_TOKEN") | |
| if not hf_token: | |
| raise gr.Error("HF_TOKEN environment variable not found. Please add your Hugging Face token to the Space settings.") | |
| return hf_token | |
| def query_api(image_bytes, prompt, seed, guidance_scale, steps, progress_callback=None): | |
| """Send request to the API using HF Router for fal.ai provider""" | |
| import requests | |
| import json | |
| import base64 | |
| hf_token = load_client() | |
| if progress_callback: | |
| progress_callback(0.1, "Submitting request...") | |
| # Use the HF router to access fal.ai provider | |
| url = "https://router.huggingface.co/fal-ai/fal-ai/flux-kontext/dev" | |
| headers = { | |
| "Authorization": f"Bearer {hf_token}", | |
| "X-HF-Bill-To": "huggingface", | |
| "Content-Type": "application/json" | |
| } | |
| # Convert image to base64 | |
| image_base64 = base64.b64encode(image_bytes).decode('utf-8') | |
| # Fixed payload structure - prompt should be at the top level | |
| payload = { | |
| "prompt": prompt, | |
| "inputs": image_base64, | |
| "seed": seed, | |
| "guidance_scale": guidance_scale, | |
| "num_inference_steps": steps | |
| } | |
| if progress_callback: | |
| progress_callback(0.3, "Processing request...") | |
| try: | |
| response = requests.post(url, headers=headers, json=payload, timeout=300) | |
| if response.status_code != 200: | |
| raise gr.Error(f"API request failed with status {response.status_code}: {response.text}") | |
| # Check if response is image bytes or JSON | |
| content_type = response.headers.get('content-type', '').lower() | |
| print(f"Response content type: {content_type}") | |
| print(f"Response length: {len(response.content)}") | |
| if 'image/' in content_type: | |
| # Direct image response | |
| if progress_callback: | |
| progress_callback(1.0, "Complete!") | |
| return response.content | |
| elif 'application/json' in content_type: | |
| # JSON response - might be queue status or result | |
| try: | |
| json_response = response.json() | |
| print(f"JSON response: {json_response}") | |
| # Check if it's a queue response | |
| if json_response.get("status") == "IN_QUEUE": | |
| if progress_callback: | |
| progress_callback(0.4, "Request queued, please wait...") | |
| raise gr.Error("Request is being processed. Please try again in a few moments.") | |
| # Handle immediate completion or result | |
| if 'images' in json_response and len(json_response['images']) > 0: | |
| image_info = json_response['images'][0] | |
| if isinstance(image_info, dict) and 'url' in image_info: | |
| # Download image from URL | |
| if progress_callback: | |
| progress_callback(0.9, "Downloading result...") | |
| img_response = requests.get(image_info['url']) | |
| if img_response.status_code == 200: | |
| if progress_callback: | |
| progress_callback(1.0, "Complete!") | |
| return img_response.content | |
| else: | |
| raise gr.Error(f"Failed to download image: {img_response.status_code}") | |
| elif isinstance(image_info, str): | |
| # Base64 encoded image | |
| if progress_callback: | |
| progress_callback(1.0, "Complete!") | |
| return base64.b64decode(image_info) | |
| elif 'image' in json_response: | |
| # Single image field | |
| if progress_callback: | |
| progress_callback(1.0, "Complete!") | |
| return base64.b64decode(json_response['image']) | |
| else: | |
| raise gr.Error(f"Unexpected JSON response format: {json_response}") | |
| except json.JSONDecodeError as e: | |
| raise gr.Error(f"Failed to parse JSON response: {str(e)}") | |
| else: | |
| # Try to treat as image bytes | |
| if len(response.content) > 1000: # Likely an image | |
| if progress_callback: | |
| progress_callback(1.0, "Complete!") | |
| return response.content | |
| else: | |
| # Small response, probably an error | |
| try: | |
| error_text = response.content.decode('utf-8') | |
| raise gr.Error(f"Unexpected response: {error_text[:500]}") | |
| except: | |
| raise gr.Error(f"Unexpected response format. Content length: {len(response.content)}") | |
| except requests.exceptions.Timeout: | |
| raise gr.Error("Request timed out. Please try again.") | |
| except requests.exceptions.RequestException as e: | |
| raise gr.Error(f"Request failed: {str(e)}") | |
| except gr.Error: | |
| # Re-raise Gradio errors as-is | |
| raise | |
| except Exception as e: | |
| raise gr.Error(f"Unexpected error: {str(e)}") | |
| # --- Core Inference Function for ChatInterface --- | |
| def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps, progress=gr.Progress()): | |
| """ | |
| Performs image generation or editing based on user input from the chat interface. | |
| """ | |
| prompt = message["text"] | |
| files = message["files"] | |
| if not prompt and not files: | |
| raise gr.Error("Please provide a prompt and/or upload an image.") | |
| if randomize_seed: | |
| seed = random.randint(0, MAX_SEED) | |
| if files: | |
| print(f"Received image: {files[0]}") | |
| try: | |
| # Try to open and convert the image | |
| input_image = Image.open(files[0]) | |
| # Convert to RGB if needed (handles RGBA, P, etc.) | |
| if input_image.mode != "RGB": | |
| input_image = input_image.convert("RGB") | |
| # Auto-orient the image based on EXIF data | |
| input_image = ImageOps.exif_transpose(input_image) | |
| # Convert PIL image to bytes | |
| img_byte_arr = io.BytesIO() | |
| input_image.save(img_byte_arr, format='PNG') | |
| img_byte_arr.seek(0) | |
| image_bytes = img_byte_arr.getvalue() | |
| except Exception as e: | |
| raise gr.Error(f"Could not process the uploaded image: {str(e)}. Please try uploading a different image format (JPEG, PNG, WebP).") | |
| progress(0.1, desc="Processing image...") | |
| else: | |
| # For text-to-image, we need a placeholder image or handle differently | |
| # FLUX.1 Kontext is primarily an image-to-image model | |
| raise gr.Error("This model (FLUX.1 Kontext) requires an input image. Please upload an image to edit.") | |
| try: | |
| # Make API request | |
| result_bytes = query_api(image_bytes, prompt, seed, guidance_scale, steps, progress_callback=progress) | |
| # Try to convert response bytes to PIL Image | |
| try: | |
| image = Image.open(io.BytesIO(result_bytes)) | |
| except Exception as img_error: | |
| print(f"Failed to open image: {img_error}") | |
| print(f"Image bytes type: {type(result_bytes)}, length: {len(result_bytes) if hasattr(result_bytes, '__len__') else 'unknown'}") | |
| # Try to decode as base64 if direct opening failed | |
| try: | |
| import base64 | |
| decoded_bytes = base64.b64decode(result_bytes) | |
| image = Image.open(io.BytesIO(decoded_bytes)) | |
| except: | |
| raise gr.Error(f"Could not process API response as image. Response length: {len(result_bytes) if hasattr(result_bytes, '__len__') else 'unknown'}") | |
| progress(1.0, desc="Complete!") | |
| return gr.Image(value=image) | |
| except gr.Error: | |
| # Re-raise gradio errors as-is | |
| raise | |
| except Exception as e: | |
| raise gr.Error(f"Failed to generate image: {str(e)}") | |
| # --- UI Definition using gr.ChatInterface --- | |
| seed_slider = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42) | |
| randomize_checkbox = gr.Checkbox(label="Randomize seed", value=False) | |
| guidance_slider = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=2.5) | |
| steps_slider = gr.Slider(label="Steps", minimum=1, maximum=30, value=28, step=1) | |
| demo = gr.ChatInterface( | |
| fn=chat_fn, | |
| title="FLUX.1 Kontext [dev] - HF Inference Client", | |
| description="""<p style='text-align: center;'> | |
| A simple chat UI for the <b>FLUX.1 Kontext [dev]</b> model using Hugging Face Inference Client approach. | |
| <br> | |
| <b>Upload an image</b> and type your editing instructions (e.g., "Turn the cat into a tiger", "Add a hat"). | |
| <br> | |
| This model specializes in understanding context and making precise edits to your images. | |
| <br> | |
| Find the model on <a href='https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev' target='_blank'>Hugging Face</a>. | |
| </p>""", | |
| multimodal=True, | |
| textbox=gr.MultimodalTextbox( | |
| file_types=["image"], | |
| placeholder="Upload an image and type your editing instructions...", | |
| render=False | |
| ), | |
| additional_inputs=[ | |
| seed_slider, | |
| randomize_checkbox, | |
| guidance_slider, | |
| steps_slider | |
| ], | |
| theme="soft" | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |