Upload folder using huggingface_hub

app.py

@@ -184,7 +184,7 @@ class GenerationManager:
             }
             payload = {
                 "prompt": prompt,
-                "enable_safety_checker": False,
+                "enable_safety_checker": True,
                 "enable_base64_output": True,  # Enable base64 output
                 "size": "1024*1024",
                 "seed": -1,
@@ -295,6 +295,7 @@ 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"):
@@ -308,7 +309,7 @@ async def poll_once(manager, backend, request_id):
                                     output, str
                             ) and not output.startswith("data:image"):
                                 # Convert raw base64 to data URI format
-                                return f"data:image/png;base64,{output}"
+                                return f"data:image/jpeg;base64,{output}"
                             else:
                                 # Already in data URI format
                                 return output
@@ -329,23 +330,45 @@ async def poll_once(manager, backend, request_id):
                 raise Exception(f"Poll error: {response.status}")
 
 
+# Store recent generations
+recent_generations = []
+
+# Example prompts
+example_prompts = [
+    "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
 
-    gr.Markdown("# 🌊 HiDream Arena powered by WaveSpeed AI Image Generator")
+    gr.Markdown("# 🌊 WaveSpeed AI HiDream Arena")
 
     # 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."
+        """
+        [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(
-        "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."
+        """
+        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
+            )
             input_text = gr.Textbox(
+                example_prompts[0],
                 label="Enter your prompt",
                 placeholder="Type here...",
                 lines=3,
@@ -353,6 +376,8 @@ 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])
+
     # Two status boxes - small (default) and big (during generation)
     small_status_box = gr.Markdown("Ready to generate images",
                                    elem_id="small-status")
@@ -374,6 +399,24 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
     performance_plot = gr.Plot(label="Performance Metrics")
 
+    with gr.Accordion("Recent Generations (last 32)", open=False):
+        recent_gallery = gr.Gallery(label="Prompt and Output")
+
+    def update_recent_gallery(prompt, results):
+        recent_generations.append({
+            "prompt": prompt,
+            "flux-dev": results["flux-dev"],
+            "hidream-dev": results["hidream-dev"],
+            "hidream-full": results["hidream-full"],
+        })
+        if len(recent_generations) > 32:
+            recent_generations.pop(0)
+        gallery_items = [
+            (r["prompt"], r["flux-dev"], r["hidream-dev"], r["hidream-full"])
+            for r in reversed(recent_generations)
+        ]
+        return gr.update(value=gallery_items)
+
     # Add custom CSS for the big status box
     css = """
     #big-status-row {
@@ -442,6 +485,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 None,
                 None,
                 session_id,  # Return the session ID
+                None,
             )
             return
 
@@ -459,6 +503,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             None,
             None,
             session_id,  # Return the session ID
+            None,
         )
 
         # For production mode:
@@ -514,6 +559,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                                     (manager.get_performance_plot()
                                      if any(completed_backends) else None),
                                     session_id,
+                                    None,
                                 )
                     except Exception as e:
                         print(f"Error polling {backend}: {str(e)}")
@@ -526,6 +572,8 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                             if len(completed_backends) == 3 else
                             "⚠️ Some generations timed out")
 
+            gallery_update = update_recent_gallery(prompt, results)
+
             # Final yield
             yield (
                 final_status,
@@ -537,6 +585,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 results["hidream-full"],
                 manager.get_performance_plot(),
                 session_id,
+                gallery_update,
             )
 
         except Exception as e:
@@ -552,6 +601,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 None,
                 None,
                 session_id,
+                None,
             )
 
     # Schedule periodic cleanup of old sessions
@@ -577,6 +627,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             best_output,
             performance_plot,
             session_id,  # Update the session ID
+            recent_gallery,  # Update the gallery
         ],
         api_name="generate",
         max_batch_size=10,  # Process up to 10 requests at once