Upload folder using huggingface_hub

app.py

@@ -9,6 +9,7 @@ from dotenv import load_dotenv
 import gradio as gr
 import random
 import torch
+import io
 from PIL import Image, ImageDraw, ImageFont
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
@@ -106,6 +107,11 @@ def image_to_base64(file_path):
         return base64.b64encode(f.read()).decode()
 
 
+def decode_base64_to_image(base64_str):
+    image_data = base64.b64decode(base64_str)
+    return Image.open(io.BytesIO(image_data))
+
+
 def generate_image(image_file,
                    prompt,
                    seed,
@@ -116,13 +122,13 @@ def generate_image(image_file,
             safety_level = classify_prompt(prompt)
             if safety_level != 0:
                 error_img = create_error_image(CLASS_NAMES[safety_level])
-                yield f"❌ Blocked: {CLASS_NAMES[safety_level]}", error_img, ""
+                yield f"❌ Blocked: {CLASS_NAMES[safety_level]}", error_img, "", None
                 return
 
         if not rate_limiter.check(session_id):
             error_img = create_error_image(
                 "Hourly limit exceeded (20 requests)")
-            yield "❌ Too many requests, please try again later", error_img, ""
+            yield "❌ Too many requests, please try again later", error_img, "", None
             return
 
         session = session_manager.get_session(session_id)
@@ -138,14 +144,15 @@ def generate_image(image_file,
             error_messages.append("Prompt cannot be empty")
         if error_messages:
             error_img = create_error_image(" | ".join(error_messages))
-            yield "❌ Input validation failed", error_img, ""
+            yield "❌ Input validation failed", error_img, "", None
             return
 
         try:
             base64_image = image_to_base64(image_file)
+            input_image = decode_base64_to_image(base64_image)
         except Exception as e:
             error_img = create_error_image(f"File processing failed: {str(e)}")
-            yield "❌ File processing failed", error_img, ""
+            yield "❌ File processing failed", error_img, "", None
             return
 
         headers = {
@@ -180,20 +187,20 @@ def generate_image(image_file,
 
             if status == "completed":
                 elapsed = time.time() - start_time
-                image_url = data["outputs"][0]
-                session["history"].append(image_url)
-                yield f"🎉 Generation successful! Time taken {elapsed:.1f}s", image_url, image_url
+                output_url = data["outputs"][0]
+                session["history"].append(output_url)
+                yield f"🎉 Generation successful! Time taken {elapsed:.1f}s", output_url, output_url, update_recent_gallery(prompt, input_image, output_url)
                 return
             elif status == "failed":
                 raise Exception(data.get("error", "Unknown error"))
             else:
-                yield f"⏳ Current status: {status.capitalize()}...", None, None
+                yield f"⏳ Current status: {status.capitalize()}...", None, None, None
 
         raise Exception("Generation timed out")
 
     except Exception as e:
         error_img = create_error_image(str(e))
-        yield f"❌ Generation failed: {str(e)}", error_img, ""
+        yield f"❌ Generation failed: {str(e)}", error_img, "", None
 
 
 def cleanup_task():
@@ -202,6 +209,9 @@ def cleanup_task():
         time.sleep(3600)
 
 
+# Store recent generations
+recent_generations = []
+
 with gr.Blocks(theme=gr.themes.Soft(),
                css="""
     .status-box { padding: 10px; border-radius: 5px; margin: 5px; }
@@ -212,27 +222,29 @@ with gr.Blocks(theme=gr.themes.Soft(),
 
     session_id = gr.State(str(uuid.uuid4()))
 
-    gr.Markdown("# 🖼️FLUX Kontext Dev Ultra Fast Live On Wavespeed AI")
+    gr.Markdown("# 🖼️FLUX Kontext Dev Ultra Fast Live")
     gr.Markdown(
         "FLUX Kontext Dev is a new SOTA image editing model published by Black Forest Labs. We have deployed it on [WaveSpeedAI](https://wavespeed.ai/) for ultra-fast image editing. You can use it to edit images in various styles, add objects, or even change the mood of the image. It supports both text prompts and image inputs."
     )
     gr.Markdown(
-        "[FLUX Kontext Dev on WaveSpeedAI](https://wavespeed.ai/models/wavespeed-ai/flux-kontext-dev)"
+        "- [FLUX Kontext Dev on WaveSpeedAI](https://wavespeed.ai/models/wavespeed-ai/flux-kontext-dev)"
+        "- [FLUX Kontext Dev LoRA on WaveSpeedAI](https://wavespeed.ai/models/wavespeed-ai/flux-kontext-dev-lora)"
+        "- [FLUX Kontext Dev Ultra Fast on WaveSpeedAI](https://wavespeed.ai/models/wavespeed-ai/flux-kontext-dev-ultra-fast)"
+        "- [FLUX Kontext Dev LoRA Ultra Fast on WaveSpeedAI](https://wavespeed.ai/models/wavespeed-ai/flux-kontext-dev-lora-ultra-fast)"
     )
     gr.Markdown(
-        "[FLUX Kontext Dev Ultra Fast on WaveSpeedAI](https://wavespeed.ai/models/wavespeed-ai/flux-kontext-dev-ultra-fast)"
     )
 
     with gr.Row():
         with gr.Column(scale=1):
+            prompt = gr.Textbox(label="Prompt",
+                                placeholder="Please enter your prompt...",
+                                lines=3)
             image_file = gr.Image(label="Upload Image",
                                   type="filepath",
                                   sources=["upload"],
                                   interactive=True,
                                   image_mode="RGB")
-            prompt = gr.Textbox(label="Prompt",
-                                placeholder="Please enter your prompt...",
-                                lines=3)
             seed = gr.Number(label="seed",
                              value=-1,
                              minimum=-1,
@@ -243,11 +255,11 @@ with gr.Blocks(theme=gr.themes.Soft(),
                                         value=True,
                                         interactive=False)
         with gr.Column(scale=1):
+            status = gr.Textbox(label="Status", elem_classes=["status-box"])
             output_image = gr.Image(label="Generated Result")
             output_url = gr.Textbox(label="Image URL",
                                     interactive=True,
                                     visible=False)
-            status = gr.Textbox(label="Status", elem_classes=["status-box"])
             submit_btn = gr.Button("Start Generation", variant="primary")
     gr.Examples(
         examples=[
@@ -275,12 +287,36 @@ with gr.Blocks(theme=gr.themes.Soft(),
         inputs=[prompt, image_file],
         label="Examples")
 
+    with gr.Accordion("Recent Generations (last 16)", open=False):
+        recent_gallery = gr.Gallery(label="Prompt and Output",
+                                    columns=3,
+                                    interactive=False)
+
+    def get_recent_gallery_items():
+        gallery_items = []
+        for r in reversed(recent_generations):
+            if any(x is None for x in r.values()):
+                continue
+            gallery_items.append((r["input"], f"Input: {r['prompt']}"))
+            gallery_items.append((r["output"], f"Output: {r['prompt']}"))
+        return gr.update(value=gallery_items)
+
+    def update_recent_gallery(prompt, input_image, output_image):
+        recent_generations.append({
+            "prompt": prompt,
+            "input": input_image,
+            "output": output_image,
+        })
+        if len(recent_generations) > 16:
+            recent_generations.pop(0)
+        return get_recent_gallery_items()
+
     random_btn.click(fn=lambda: random.randint(0, 999999), outputs=seed)
 
     submit_btn.click(
         generate_image,
         inputs=[image_file, prompt, seed, session_id, enable_safety],
-        outputs=[status, output_image, output_url],
+        outputs=[status, output_image, output_url, recent_gallery],
         api_name=False,
     )