Update app.py

app.py

@@ -10,20 +10,14 @@ from PIL import Image
 BTEN_API_KEY = os.getenv("API_KEY")
 URL = os.getenv("URL")
 
+
 def image_to_base64(image: Image.Image) -> str:
-    """Convert a PIL image to a base64-encoded PNG string."""
     with io.BytesIO() as buffer:
         image.save(buffer, format="PNG")
         return base64.b64encode(buffer.getvalue()).decode("utf-8")
 
 
 def ensure_image(img) -> Image.Image:
-    """
-    Ensure the input is a PIL Image.
-    If it's already a PIL Image, return it.
-    If it's a string (file path), open it.
-    If it's a dict with a "name" key, open the file at that path.
-    """
     if isinstance(img, Image.Image):
         return img
     elif isinstance(img, str):
@@ -34,19 +28,7 @@ def ensure_image(img) -> Image.Image:
         raise ValueError("Cannot convert input to a PIL Image.")
 
 
-def call_baseten_generate(
-    image: Image.Image,
-    prompt: str,
-    steps: int,
-    strength: float,
-    height: int,
-    width: int,
-    lora_name: str,
-    remove_bg: bool,
-) -> Image.Image | None:
-    """
-    Call the Baseten /predict endpoint with provided parameters and return the generated image.
-    """
+def call_baseten_generate(image: Image.Image, prompt: str, steps: int, strength: float, height: int, width: int, lora_name: str, remove_bg: bool) -> Image.Image | None:
     image = ensure_image(image)
     b64_image = image_to_base64(image)
     payload = {
@@ -59,14 +41,10 @@ def call_baseten_generate(
         "lora_name": lora_name,
         "bgrm": remove_bg,
     }
-    if not BTEN_API_KEY:
-        headers = {"Authorization": f"Api-Key {os.getenv('API_KEY')}"}
-    else:
-        headers = {"Authorization": f"Api-Key {BTEN_API_KEY}"}
+    headers = {"Authorization": f"Api-Key {BTEN_API_KEY or os.getenv('API_KEY')}"}
     try:
         if not URL:
             raise ValueError("The URL environment variable is not set.")
-
         response = requests.post(URL, headers=headers, json=payload)
         if response.status_code == 200:
             data = response.json()
@@ -83,20 +61,17 @@ def call_baseten_generate(
         return None
 
 
-# Mode defaults for each tab.
+# ================== MODE CONFIG =====================
 
-Mode = TypedDict(
-    "Mode",
-    {
-        "model": str,
-        "prompt": str,
-        "default_strength": float,
-        "default_height": int,
-        "default_width": int,
-        "models": list[str],
-        "remove_bg": bool,
-    },
-)
+Mode = TypedDict("Mode", {
+    "model": str,
+    "prompt": str,
+    "default_strength": float,
+    "default_height": int,
+    "default_width": int,
+    "models": list[str],
+    "remove_bg": bool,
+})
 
 MODE_DEFAULTS: dict[str, Mode] = {
     "Subject Generation": {
@@ -105,12 +80,7 @@ MODE_DEFAULTS: dict[str, Mode] = {
         "default_strength": 1.2,
         "default_height": 512,
         "default_width": 512,
-        "models": [
-            "zendsd_512_146000",
-            "subject_99000_512",
-            # "zen_pers_11000",
-            "zen_26000_512",
-        ],
+        "models": ["zendsd_512_146000", "subject_99000_512", "zen_26000_512"],
         "remove_bg": True,
     },
     "Background Generation": {
@@ -119,20 +89,7 @@ MODE_DEFAULTS: dict[str, Mode] = {
         "default_strength": 1.2,
         "default_height": 1024,
         "default_width": 1024,
-        "models": [
-            "bgwlight_15000_1024",
-            # "rmgb_12000_1024",
-            "bg_canny_58000_1024",
-            # "gen_back_3000_1024",
-            "gen_back_7000_1024",
-            # "gen_bckgnd_18000_512",
-            # "gen_bckgnd_18000_512",
-            # "loose_25000_512",
-            # "looser_23000_1024",
-            # "looser_bg_gen_21000_1280",
-            # "old_looser_46000_1024",
-            # "relight_bg_gen_31000_1024",
-        ],
+        "models": ["bgwlight_15000_1024", "bg_canny_58000_1024", "gen_back_7000_1024"],
         "remove_bg": True,
     },
     "Canny": {
@@ -150,9 +107,7 @@ MODE_DEFAULTS: dict[str, Mode] = {
         "default_strength": 1.2,
         "default_height": 1024,
         "default_width": 1024,
-        "models": [
-            "depth_9800_1024",
-        ],
+        "models": ["depth_9800_1024"],
         "remove_bg": True,
     },
     "Deblurring": {
@@ -161,63 +116,87 @@ MODE_DEFAULTS: dict[str, Mode] = {
         "default_strength": 1.2,
         "default_height": 1024,
         "default_width": 1024,
-        "models": ["deblurr_1024_10000"],  # "slight_deblurr_18000",
+        "models": ["deblurr_1024_10000"],
         "remove_bg": False,
     },
 }
 
+# ================== PRESET EXAMPLES =====================
+
+MODE_EXAMPLES = {
+    "Subject Generation": [
+        ["assets/subj1.jpg", "Close-up portrait of a fruit bowl", "assets/subj1_out.jpg"],
+        ["assets/subj2.jpg", "A penguin standing in snow", "assets/subj2_out.jpg"],
+        ["assets/subj3.jpg", "A cat with glowing eyes", "assets/subj3_out.jpg"],
+        ["assets/subj4.jpg", "A child playing with bubbles", "assets/subj4_out.jpg"],
+        ["assets/subj5.jpg", "A stylish young man in neon lights", "assets/subj5_out.jpg"],
+        ["assets/subj6.jpg", "Old man with a mysterious look", "assets/subj6_out.jpg"],
+    ],
+
+    "Background Generation": [
+        ["assets/bg1.jpg", "Modern living room with plants", "assets/bg1_out.jpg"],
+        ["assets/bg2.jpg", "Fantasy forest background", "assets/bg2_out.jpg"],
+        ["assets/bg3.jpg", "Futuristic cityscape", "assets/bg3_out.jpg"],
+        ["assets/bg4.jpg", "Minimalist white studio", "assets/bg4_out.jpg"],
+        ["assets/bg5.jpg", "Snowy mountain landscape", "assets/bg5_out.jpg"],
+        ["assets/bg6.jpg", "Golden sunset over the sea", "assets/bg6_out.jpg"],
+    ],
+
+    "Canny": [
+        ["assets/canny1.jpg", "A neon cyberpunk city skyline", "assets/canny1_out.jpg"],
+        ["assets/canny2.jpg", "A robot walking in the fog", "assets/canny2_out.jpg"],
+        ["assets/canny3.jpg", "A futuristic vehicle parked under a bridge", "assets/canny3_out.jpg"],
+        ["assets/canny4.jpg", "Sci-fi lab interior with glowing machinery", "assets/canny4_out.jpg"],
+        ["assets/canny5.jpg", "A portrait of a woman outlined in neon", "assets/canny5_out.jpg"],
+        ["assets/canny6.jpg", "Post-apocalyptic abandoned street", "assets/canny6_out.jpg"],
+    ],
+
+    "Depth": [
+        ["assets/depth1.jpg", "A narrow alleyway with deep perspective", "assets/depth1_out.jpg"],
+        ["assets/depth2.jpg", "A mountain road vanishing into the distance", "assets/depth2_out.jpg"],
+        ["assets/depth3.jpg", "A hallway with strong depth of field", "assets/depth3_out.jpg"],
+        ["assets/depth4.jpg", "A misty forest path stretching far away", "assets/depth4_out.jpg"],
+        ["assets/depth5.jpg", "A bridge over a deep canyon", "assets/depth5_out.jpg"],
+        ["assets/depth6.jpg", "An underground tunnel with receding arches", "assets/depth6_out.jpg"],
+    ],
+
+    "Deblurring": [
+        ["assets/deblur1.jpg", "", "assets/deblur1_out.jpg"],
+        ["assets/deblur2.jpg", "", "assets/deblur2_out.jpg"],
+        ["assets/deblur3.jpg", "", "assets/deblur3_out.jpg"],
+        ["assets/deblur4.jpg", "", "assets/deblur4_out.jpg"],
+        ["assets/deblur5.jpg", "", "assets/deblur5_out.jpg"],
+        ["assets/deblur6.jpg", "", "assets/deblur6_out.jpg"],
+    ],
+}
+
+
+# ================== UI =====================
 
 header = """
 <h1>🌍 ZenCtrl / FLUX</h1>
 <div align="center" style="line-height: 1;">
-    <a href="https://github.com/FotographerAI/ZenCtrl/tree/main" target="_blank" style="margin: 2px;" name="github_repo_link"><img src="https://img.shields.io/badge/GitHub-Repo-181717.svg" alt="GitHub Repo" style="display: inline-block; vertical-align: middle;"></a>
-    <a href="https://huggingface.co/spaces/fotographerai/ZenCtrl" target="_blank" name="huggingface_space_link"><img src="https://img.shields.io/badge/🤗_HuggingFace-Space-ffbd45.svg" alt="HuggingFace Space" style="display: inline-block; vertical-align: middle;"></a>
-    <a href="https://discord.com/invite/b9RuYQ3F8k" target="_blank" style="margin: 2px;" name="discord_link"><img src="https://img.shields.io/badge/Discord-Join-7289da.svg?logo=discord" alt="Discord" style="display: inline-block; vertical-align: middle;"></a>
-    <a href="https://fotographer.ai/" target="_blank" style="margin: 2px;" name="lp_link"><img src="https://img.shields.io/badge/Website-Landing_Page-blue" alt="LP" style="display: inline-block; vertical-align: middle;"></a>
-    <a href="https://x.com/FotographerAI" target="_blank" style="margin: 2px;" name="twitter_link"><img src="https://img.shields.io/twitter/follow/FotographerAI?style=social" alt="X" style="display: inline-block; vertical-align: middle;"></a>
+    <a href="https://github.com/FotographerAI/ZenCtrl/tree/main" target="_blank"><img src="https://img.shields.io/badge/GitHub-Repo-181717.svg"></a>
+    <a href="https://huggingface.co/spaces/fotographerai/ZenCtrl" target="_blank"><img src="https://img.shields.io/badge/🤗_HuggingFace-Space-ffbd45.svg"></a>
+    <a href="https://discord.com/invite/b9RuYQ3F8k" target="_blank"><img src="https://img.shields.io/badge/Discord-Join-7289da.svg?logo=discord"></a>
 </div>
 """
 
-defaults = MODE_DEFAULTS["Subject Generation"]
-
-
 with gr.Blocks(title="🌍 ZenCtrl") as demo:
     gr.HTML(header)
-    gr.Markdown(
-        """
-        # ZenCtrl Demo
-        [WIP] One Agent to Generate multi-view, diverse-scene, and task-specific high-resolution images from a single subject image—without fine-tuning.
-        We are first releasing some of the task specific weights and will release the codes soon.
-        The goal is to unify all of the visual content generation tasks with a single LLM...
-        
-        **Modes:**
-        - **Subject Generation:** Focuses on generating detailed subject portraits.
-        - **Background Generation:** Creates dynamic, vibrant backgrounds:
-            You can generate part of the image from sketch while keeping part of it as it is.
-        - **Canny:** Emphasizes strong edge detection.
-        - **Depth:** Produces images with realistic depth and perspective.
-        
-        For more details, shoot us a message on discord.
-        """
-    )
+    gr.Markdown("# ZenCtrl Demo")
+
     with gr.Tabs():
         for mode in MODE_DEFAULTS:
             with gr.Tab(mode):
                 defaults = MODE_DEFAULTS[mode]
                 gr.Markdown(f"### {mode} Mode")
-                gr.Markdown(f"**Default Model:** {defaults['model']}")
 
                 with gr.Row():
-                    with gr.Column(scale=2, min_width=370):
-                        input_image = gr.Image(
-                            label="Upload Image",
-                            type="pil",
-                            scale=3,
-                            height=370,
-                            min_width=100,
-                        )
+                    with gr.Column(scale=2):
+                        input_image = gr.Image(label="Input Image", type="pil")
                         generate_button = gr.Button("Generate")
-                        with gr.Blocks(title="Options"):
+                        with gr.Blocks():
                             model_dropdown = gr.Dropdown(
                                 label="Model",
                                 choices=defaults["models"],
@@ -229,20 +208,12 @@ with gr.Blocks(title="🌍 ZenCtrl") as demo:
                             )
 
                     with gr.Column(scale=2):
-                        output_image = gr.Image(
-                            label="Generated Image",
-                            type="pil",
-                            height=573,
-                            scale=4,
-                            min_width=100,
-                        )
+                        output_image = gr.Image(label="Generated Image", type="pil")
 
-                gr.Markdown("#### Prompt")
                 prompt_box = gr.Textbox(
                     label="Prompt", value=defaults["prompt"], lines=2
                 )
 
-                # Wrap generation parameters in an Accordion for collapsible view.
                 with gr.Accordion("Generation Parameters", open=False):
                     with gr.Row():
                         step_slider = gr.Slider(
@@ -271,26 +242,8 @@ with gr.Blocks(title="🌍 ZenCtrl") as demo:
                             label="Width",
                         )
 
-                def on_generate_click(
-                    model_name,
-                    prompt,
-                    steps,
-                    strength,
-                    height,
-                    width,
-                    remove_bg,
-                    image,
-                ):
-                    return call_baseten_generate(
-                        image,
-                        prompt,
-                        steps,
-                        strength,
-                        height,
-                        width,
-                        model_name,
-                        remove_bg,
-                    )
+                def on_generate_click(model_name, prompt, steps, strength, height, width, remove_bg, image):
+                    return call_baseten_generate(image, prompt, steps, strength, height, width, model_name, remove_bg)
 
                 generate_button.click(
                     fn=on_generate_click,
@@ -305,9 +258,16 @@ with gr.Blocks(title="🌍 ZenCtrl") as demo:
                         input_image,
                     ],
                     outputs=[output_image],
-                    concurrency_limit=None
                 )
 
+                # ---------------- Templates --------------------
+                gr.Dataset(
+                    label="Presets (Input / Prompt / Output)",
+                    headers=["Input", "Prompt", "Output"],
+                    components=[input_image, prompt_box, output_image],
+                    samples=MODE_EXAMPLES.get(mode, []),
+                    samples_per_page=6,
+                )
 
 if __name__ == "__main__":
     demo.launch()