Upload folder using huggingface_hub

README.md

@@ -1,12 +1,14 @@
 ---
-title: Hidream Arena
-emoji: 📚
-colorFrom: pink
-colorTo: green
+title: HiDream Arena
+emoji: 🖼
+colorFrom: purple
+colorTo: red
 sdk: gradio
-sdk_version: 5.25.0
+sdk_version: 5.0.1
 app_file: app.py
 pinned: false
+license: mit
+short_description: Arena for HiDream-I1-dev / -full and FLUX-dev
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,592 @@
+import gradio as gr
+import asyncio
+import aiohttp
+import time
+from datetime import datetime
+import plotly.graph_objects as go
+from typing import Dict, List
+import os
+from dotenv import load_dotenv
+import json
+from PIL import Image, ImageDraw, ImageFont
+import uuid
+import threading
+
+# Load environment variables first
+load_dotenv()
+
+# Constants
+API_BASE_URL = "https://api.wavespeed.ai/api/v2"
+API_KEY = os.getenv("WAVESPEED_API_KEY")  # Move API_KEY to global scope
+
+if not API_KEY:
+    raise ValueError("WAVESPEED_API_KEY not found in environment variables")
+
+# Rest of constants
+BACKENDS = {
+    "flux-dev": {
+        "endpoint": f"{API_BASE_URL}/wavespeed-ai/flux-dev-ultra-fast",
+        "name": "Flux-dev",
+        "color": "#FF9800",
+    },
+    "hidream-dev": {
+        "endpoint": f"{API_BASE_URL}/wavespeed-ai/hidream-i1-dev",
+        "name": "HiDream-dev",
+        "color": "#2196F3",
+    },
+    "hidream-full": {
+        "endpoint": f"{API_BASE_URL}/wavespeed-ai/hidream-i1-full",
+        "name": "HiDream-full",
+        "color": "#4CAF50",
+    },
+}
+
+
+class BackendStatus:
+
+    def __init__(self):
+        self.reset()
+        self.history: List[Dict] = []
+
+    def reset(self):
+        self.status = "idle"
+        self.progress = 0
+        self.start_time = None
+        self.end_time = None
+
+    def start(self):
+        self.status = "processing"
+        self.progress = 0
+        self.start_time = time.time()
+        self.end_time = None
+
+    def complete(self):
+        self.status = "completed"
+        self.progress = 100
+        self.end_time = time.time()
+        self.history.append({
+            "timestamp": datetime.now(),
+            "duration": self.end_time - self.start_time
+        })
+
+    def fail(self):
+        self.status = "failed"
+        self.end_time = time.time()
+
+
+class SessionManager:
+    _instances = {}
+    _lock = threading.Lock()
+
+    @classmethod
+    def get_manager(cls, session_id=None):
+        if session_id is None:
+            session_id = str(uuid.uuid4())
+
+        with cls._lock:
+            if session_id not in cls._instances:
+                cls._instances[session_id] = GenerationManager()
+            return session_id, cls._instances[session_id]
+
+    @classmethod
+    def cleanup_old_sessions(cls, max_age=3600):  # 1 hour default
+        current_time = time.time()
+        with cls._lock:
+            to_remove = []
+            for session_id, manager in cls._instances.items():
+                if (hasattr(manager, "last_activity")
+                        and current_time - manager.last_activity > max_age):
+                    to_remove.append(session_id)
+
+            for session_id in to_remove:
+                del cls._instances[session_id]
+
+
+class GenerationManager:
+
+    def __init__(self):
+        self.backend_statuses = {
+            backend: BackendStatus()
+            for backend in BACKENDS
+        }
+        self.last_activity = time.time()
+
+    def update_activity(self):
+        self.last_activity = time.time()
+
+    def get_performance_plot(self):
+        fig = go.Figure()
+
+        has_data = False
+
+        for backend, status in self.backend_statuses.items():
+            durations = [h["duration"] for h in status.history]
+            if durations:
+                has_data = True
+                avg_duration = sum(durations) / len(durations)
+                # Use bar chart instead of box plot
+                fig.add_trace(
+                    go.Bar(
+                        y=[avg_duration],  # Average duration
+                        x=[BACKENDS[backend]["name"]],  # Backend name
+                        name=BACKENDS[backend]["name"],
+                        marker_color=BACKENDS[backend]["color"],
+                        text=[f"{avg_duration:.2f}s"],  # Show time in seconds
+                        textposition="auto",
+                        width=[0.5],  # Make bars narrower
+                    ))
+
+        # Set a minimum y-axis range if we have data
+        if has_data:
+            max_duration = max([
+                max([h["duration"] for h in status.history] or [0])
+                for status in self.backend_statuses.values()
+            ])
+            # Add 20% padding to the top
+            y_max = max_duration * 1.2
+            # Ensure the y-axis always starts at 0
+            fig.update_yaxes(range=[0, y_max])
+
+        fig.update_layout(
+            title="Average Generation Time",
+            yaxis_title="Seconds",
+            xaxis_title="",
+            showlegend=False,
+            template="simple_white",
+            height=400,  # Increase height
+            margin=dict(l=50, r=50, t=50, b=50),  # Add margins
+            font=dict(size=14),  # Larger font
+        )
+
+        # Make sure we have a valid figure even if no data
+        if not has_data:
+            fig.add_annotation(
+                text="No timing data available yet",
+                xref="paper",
+                yref="paper",
+                x=0.5,
+                y=0.5,
+                showarrow=False,
+                font=dict(size=16),
+            )
+
+        return fig
+
+    async def submit_task(self, backend: str, prompt: str) -> str:
+        status = self.backend_statuses[backend]
+        status.start()
+
+        try:
+            url = BACKENDS[backend]["endpoint"]
+            headers = {
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {API_KEY}",
+            }
+            payload = {
+                "prompt": prompt,
+                "enable_safety_checker": False,
+                "enable_base64_output": True,  # Enable base64 output
+                "size": "1024*1024",
+                "seed": -1,
+            }
+
+            if backend == "flux-dev":
+                payload.update({
+                    "guidance_scale": 3.5,
+                    "num_images": 1,
+                    "num_inference_steps": 28,
+                    "strength": 0.8,
+                })
+
+            print(f"Submitting task to {backend}")
+            print(f"URL: {url}")
+            print(f"Payload: {json.dumps(payload, indent=2)}")
+
+            # Use aiohttp instead of requests for async
+            async with aiohttp.ClientSession() as session:
+                async with session.post(url, headers=headers,
+                                        json=payload) as response:
+                    if response.status == 200:
+                        result = await response.json()
+                        request_id = result["data"]["id"]
+                        print(
+                            f"Task submitted successfully. Request ID: {request_id}"
+                        )
+                        return request_id
+                    else:
+                        text = await response.text()
+                        raise Exception(
+                            f"API error: {response.status}, {text}")
+
+        except Exception as e:
+            status.fail()
+            raise Exception(f"Failed to submit task: {str(e)}")
+
+    # Add this method to reset history
+    def reset_history(self):
+        """Reset history for all backends"""
+        for status in self.backend_statuses.values():
+            status.history = []  # Clear history data
+        return self
+
+
+# Helper function to create error images as data URIs
+def create_error_image(backend, error_message):
+    try:
+        import base64
+        from io import BytesIO
+
+        # Create an in-memory image
+        img = Image.new("RGB", (512, 512), color="#ffdddd")
+        draw = ImageDraw.Draw(img)
+        try:
+            font = ImageFont.truetype("Arial", 20)
+        except:
+            font = ImageFont.load_default()
+
+        # Wrap and draw error message
+        words = error_message.split(" ")
+        lines = []
+        line = ""
+        for word in words:
+            if len(line + word) < 40:
+                line += word + " "
+            else:
+                lines.append(line)
+                line = word + " "
+        if line:
+            lines.append(line)
+
+        y_position = 100
+        for line in lines:
+            draw.text((50, y_position), line, fill="black", font=font)
+            y_position += 30
+
+        # Save to a BytesIO object instead of a file
+        buffer = BytesIO()
+        img.save(buffer, format="PNG")
+        img_bytes = buffer.getvalue()
+
+        # Convert to base64 and return as data URI
+        return f"data:image/png;base64,{base64.b64encode(img_bytes).decode('utf-8')}"
+    except Exception as e:
+        print(f"Failed to create error image: {e}")
+        # Return a simple error message as fallback
+        return "Error: " + error_message
+
+
+# Fix the poll_once function to accept a manager parameter
+async def poll_once(manager, backend, request_id):
+    """Poll once and return result if complete, otherwise None"""
+    headers = {"Authorization": f"Bearer {API_KEY}"}
+    url = f"{API_BASE_URL}/predictions/{request_id}/result"
+
+    async with aiohttp.ClientSession() as session:
+        async with session.get(url, headers=headers) as response:
+            if response.status == 200:
+                result = await response.json()
+                data = result["data"]
+                current_status = data["status"]
+
+                if current_status == "completed":
+                    # IMPORTANT: Update status BEFORE returning - using the passed manager
+                    manager.backend_statuses[backend].complete()
+                    manager.update_activity()
+
+                    # Handle base64 output
+                    output = data["outputs"][0]
+
+                    # Check if it's a base64 string or URL
+                    if isinstance(output, str) and output.startswith("http"):
+                        # It's a URL - return as is
+                        return output
+                    else:
+                        # It's base64 data - format it as a data URI if needed
+                        try:
+                            # Format as data URI for Gradio to display directly
+                            if isinstance(
+                                    output, str
+                            ) and not output.startswith("data:image"):
+                                # Convert raw base64 to data URI format
+                                return f"data:image/png;base64,{output}"
+                            else:
+                                # Already in data URI format
+                                return output
+                        except Exception as e:
+                            print(f"Error processing base64 image: {e}")
+                            raise Exception(
+                                f"Failed to process base64 image: {str(e)}")
+
+                elif current_status == "failed":
+                    manager.backend_statuses[backend].fail()
+                    manager.update_activity()
+                    error = data.get("error", "Unknown error")
+                    raise Exception(error)
+
+                # Still processing
+                return None
+            else:
+                raise Exception(f"Poll error: {response.status}")
+
+
+# Use a state variable to store session ID
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    session_id = gr.State(None)  # Add this to store session ID
+
+    gr.Markdown("# 🌊 WaveSpeed AI Image Generator")
+
+    # Add the introduction with link to WaveSpeedAI
+    gr.Markdown(
+        "[WaveSpeedAI](https://wavespeed.ai/) is the global pioneer in accelerating AI-powered video and image generation."
+    )
+    gr.Markdown(
+        "Our in-house inference accelerator provides lossless speedup on image & video generation based on our rich inference optimization software stack, including our in-house inference compiler, CUDA kernel libraries and parallel computing libraries."
+    )
+
+    with gr.Row():
+        with gr.Column(scale=3):
+            input_text = gr.Textbox(
+                label="Enter your prompt",
+                placeholder="Type here...",
+                lines=3,
+            )
+        with gr.Column(scale=1):
+            generate_btn = gr.Button("Generate", variant="primary")
+
+    # Two status boxes - small (default) and big (during generation)
+    small_status_box = gr.Markdown("Ready to generate images",
+                                   elem_id="small-status")
+
+    # Big status box in its own row with styling
+    with gr.Row(elem_id="big-status-row"):
+        big_status_box = gr.Markdown("",
+                                     elem_id="big-status",
+                                     visible=False,
+                                     elem_classes="big-status-box")
+
+    with gr.Row():
+        with gr.Column():
+            draft_output = gr.Image(label="Flux-dev")
+        with gr.Column():
+            quick_output = gr.Image(label="HiDream-dev")
+        with gr.Column():
+            best_output = gr.Image(label="HiDream-full")
+
+    performance_plot = gr.Plot(label="Performance Metrics")
+
+    # Add custom CSS for the big status box
+    css = """
+    #big-status-row {
+        margin: 20px 0;
+    }
+    #big-status {
+        font-size: 28px;  /* Even larger font size */
+        font-weight: bold;
+        padding: 30px;    /* More padding */
+        background-color: #0D47A1;  /* Deeper blue background */
+        color: white;     /* White text */
+        border-radius: 10px;
+        text-align: center;
+        margin: 0 auto;
+        box-shadow: 0 6px 12px rgba(0, 0, 0, 0.2);  /* Stronger shadow */
+        animation: deep-breath 3s infinite;  /* Slower, deeper breathing animation */
+        width: 100%;      /* Full width */
+        max-width: 800px; /* Maximum width */
+        transition: all 0.3s ease;  /* Smooth transitions */
+        border-left: 6px solid #64B5F6;  /* Add a colored border */
+        border-right: 6px solid #64B5F6;  /* Add a colored border */
+    }
+
+    /* Deeper breathing animation */
+    @keyframes deep-breath {
+        0% { 
+            opacity: 0.7; 
+            transform: scale(0.98);
+            box-shadow: 0 4px 8px rgba(0, 0, 0, 0.2);
+        }
+        50% { 
+            opacity: 1; 
+            transform: scale(1.01);
+            box-shadow: 0 8px 16px rgba(0, 0, 0, 0.3);
+        }
+        100% { 
+            opacity: 0.7; 
+            transform: scale(0.98);
+            box-shadow: 0 4px 8px rgba(0, 0, 0, 0.2);
+        }
+    }
+    """
+    gr.HTML(f"<style>{css}</style>")
+
+    # Update the generation function to use session manager
+    async def generate_all_backends_with_status_boxes(prompt,
+                                                      current_session_id):
+        """Generate images with big status box during generation"""
+        # Get or create a session manager
+        session_id, manager = SessionManager.get_manager(current_session_id)
+        manager.update_activity()
+
+        # IMPORTANT: Reset history when starting a new generation
+        if prompt and prompt.strip() != "":
+            manager.reset_history()  # Clear previous performance metrics
+
+        if not prompt or prompt.strip() == "":
+            # Handle empty prompt case
+            yield (
+                "⚠️ Please enter a prompt first",
+                "⚠️ Please enter a prompt first",
+                gr.update(visible=True),
+                gr.update(visible=False),
+                None,
+                None,
+                None,
+                None,
+                session_id,  # Return the session ID
+            )
+            return
+
+        # Status message
+        status_message = f"🔄 PROCESSING: '{prompt}'"
+
+        # Initial state - clear all images, show big status box
+        yield (
+            status_message,
+            status_message,
+            gr.update(visible=True),
+            gr.update(visible=False),
+            None,
+            None,
+            None,
+            None,
+            session_id,  # Return the session ID
+        )
+
+        # For production mode:
+        completed_backends = set()
+        results = {"flux-dev": None, "hidream-dev": None, "hidream-full": None}
+
+        try:
+            # Submit all tasks
+            request_ids = {}
+            for backend in BACKENDS:
+                try:
+                    request_id = await manager.submit_task(backend, prompt)
+                    request_ids[backend] = request_id
+                except Exception as e:
+                    # Handle submission error
+                    print(f"Error submitting task for {backend}: {e}")
+                    results[backend] = create_error_image(backend, str(e))
+                    completed_backends.add(backend)
+
+            # Poll all backends until they complete
+            max_poll_attempts = 300
+            poll_attempt = 0
+
+            # Main polling loop
+            while len(completed_backends
+                      ) < 3 and poll_attempt < max_poll_attempts:
+                poll_attempt += 1
+
+                # Poll each pending backend
+                for backend in list(BACKENDS.keys()):
+                    if backend in completed_backends:
+                        continue
+
+                    try:
+                        # Only do actual API calls every few attempts to reduce load
+                        if poll_attempt % 2 == 0 or backend == "flux-dev":
+                            # Use the session manager instead of global manager
+                            result = await poll_once(manager, backend,
+                                                     request_ids[backend])
+                            if result:  # Backend completed
+                                results[backend] = result
+                                completed_backends.add(backend)
+
+                                # Yield updated state when an image completes
+                                yield (
+                                    status_message,
+                                    status_message,
+                                    gr.update(visible=True),
+                                    gr.update(visible=False),
+                                    results["flux-dev"],
+                                    results["hidream-dev"],
+                                    results["hidream-full"],
+                                    (manager.get_performance_plot()
+                                     if any(completed_backends) else None),
+                                    session_id,
+                                )
+                    except Exception as e:
+                        print(f"Error polling {backend}: {str(e)}")
+
+                # Wait between poll attempts
+                await asyncio.sleep(0.1)
+
+            # Final status
+            final_status = ("✅ All generations completed!"
+                            if len(completed_backends) == 3 else
+                            "⚠️ Some generations timed out")
+
+            # Final yield
+            yield (
+                final_status,
+                final_status,
+                gr.update(visible=False),
+                gr.update(visible=True),
+                results["flux-dev"],
+                results["hidream-dev"],
+                results["hidream-full"],
+                manager.get_performance_plot(),
+                session_id,
+            )
+
+        except Exception as e:
+            # Error handling
+            error_message = f"❌ Error: {str(e)}"
+            yield (
+                error_message,
+                error_message,
+                gr.update(visible=False),
+                gr.update(visible=True),
+                None,
+                None,
+                None,
+                None,
+                session_id,
+            )
+
+    # Schedule periodic cleanup of old sessions
+    def cleanup_task():
+        SessionManager.cleanup_old_sessions()
+        # Schedule the next cleanup
+        threading.Timer(3600, cleanup_task).start()  # Run every hour
+
+    # Start the cleanup task
+    cleanup_task()
+
+    # Update the click handler to include session_id
+    generate_btn.click(
+        fn=generate_all_backends_with_status_boxes,
+        inputs=[input_text, session_id],
+        outputs=[
+            small_status_box,
+            big_status_box,
+            big_status_box,  # visibility
+            small_status_box,  # visibility
+            draft_output,
+            quick_output,
+            best_output,
+            performance_plot,
+            session_id,  # Update the session ID
+        ],
+        api_name="generate",
+        max_batch_size=10,  # Process up to 10 requests at once
+        concurrency_limit=20,  # Allow up to 20 concurrent requests
+        concurrency_id="generation",  # Group concurrent requests under this ID
+    )
+
+# Launch with increased max_threads
+if __name__ == "__main__":
+    demo.queue(max_size=50).launch(
+        server_name="0.0.0.0",
+        max_threads=16,  # Increase thread count for better concurrency
+    )

requirements.txt

@@ -0,0 +1,4 @@
+gradio
+aiohttp
+plotly
+python-dotenv