Upload folder using huggingface_hub

app.py

@@ -264,7 +264,7 @@ def create_error_image(backend, error_message):
 
         # Save to a BytesIO object instead of a file
         buffer = BytesIO()
-        img.save(buffer, format="PNG")
+        img.save(buffer, format="JPEG")
         img_bytes = buffer.getvalue()
 
         # Convert to base64 and return as data URI
@@ -295,7 +295,6 @@ async def poll_once(manager, backend, request_id):
 
                     # Handle base64 output
                     output = data["outputs"][0]
-                    # has_nsfw_content = data["has_nsfw_contents"][0]
 
                     # Check if it's a base64 string or URL
                     if isinstance(output, str) and output.startswith("http"):
@@ -335,11 +334,12 @@ recent_generations = []
 
 # Example prompts
 example_prompts = [
-    "A deep sea diver exploring an underwater city ruins, using a palette of deep blues and silvers.",
+    "A deep sea diver exploring an underwater city ruins, using a palette of deep blues and silvers."
     "A Martian greenhouse complex that uses genetically modified crops designed to thrive in low gravity. Outside, rovers fitted with AI guidance systems patrol dusty red plains, ensuring each pressurized dome remains airtight.",
     "A sleek, futuristic sports car with glowing blue accents, racing through a virtual reality landscape, 3D render",
 ]
 
+
 # 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
@@ -347,19 +347,25 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🌊 WaveSpeed AI HiDream Arena")
 
     # Add the introduction with link to WaveSpeedAI
-    gr.Markdown("""
+    gr.Markdown(
+        """
         [WaveSpeedAI](https://wavespeed.ai/) is the global pioneer in accelerating AI-powered video and image generation.
         "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."
-        """)
-    gr.Markdown("""
+        """
+    )
+    gr.Markdown(
+        """
         This demo showcases the performance and outputs of leading image generation models, including HiDream and Flux, on our accelerated inference platform.
-        """)
+        """
+    )
 
     with gr.Row():
         with gr.Column(scale=3):
-            example_dropdown = gr.Dropdown(choices=example_prompts,
-                                           label="Choose an example prompt",
-                                           interactive=True)
+            example_dropdown = gr.Dropdown(
+                choices=example_prompts,
+                label="Choose an example prompt",
+                interactive=True
+            )
             input_text = gr.Textbox(
                 example_prompts[0],
                 label="Enter your prompt",
@@ -369,9 +375,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         with gr.Column(scale=1):
             generate_btn = gr.Button("Generate", variant="primary")
 
-    example_dropdown.change(lambda ex: ex,
-                            inputs=[example_dropdown],
-                            outputs=[input_text])
+    example_dropdown.change(lambda ex: ex, inputs=[example_dropdown], outputs=[input_text])
 
     # Two status boxes - small (default) and big (during generation)
     small_status_box = gr.Markdown("Ready to generate images",
@@ -395,10 +399,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     performance_plot = gr.Plot(label="Performance Metrics")
 
     with gr.Accordion("Recent Generations (last 8)", open=False):
-        recent_gallery = gr.Gallery(label="Recent Generations",
-                                    columns=3,
-                                    rows=1,
-                                    show_label=True)
+        recent_gallery = gr.Gallery(label="Prompt and Output")
 
     def update_recent_gallery(prompt, results):
         recent_generations.append({
@@ -409,15 +410,15 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         })
         if len(recent_generations) > 8:
             recent_generations.pop(0)
-
-        # Flatten the images from each generation
-        all_images = []
-        for gen in reversed(recent_generations):  # Show newest first
-            all_images.extend(
-                [gen["flux-dev"], gen["hidream-dev"], gen["hidream-full"]])
-
-        # Update the Gradio gallery with the new list
-        return recent_gallery.update(value=all_images)
+        gallery_items = []
+        for r in reversed(recent_generations):
+            gallery_items.append(
+                (r["flux-dev"], f"FLUX-dev: {r['prompt']}"))
+            gallery_items.append(
+                (r["hidream-dev"], f"HiDream-dev: {r['prompt']}"))
+            gallery_items.append(
+                (r["hidream-full"], f"HiDream-full: {r['prompt']}"))
+        return gr.update(value=gallery_items)
 
     # Add custom CSS for the big status box
     css = """