Update app.py

app.py

@@ -10,14 +10,20 @@ 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):
@@ -28,7 +34,19 @@ 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:
+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.
+    """
     image = ensure_image(image)
     b64_image = image_to_base64(image)
     payload = {
@@ -41,10 +59,14 @@ def call_baseten_generate(image: Image.Image, prompt: str, steps: int, strength:
         "lora_name": lora_name,
         "bgrm": remove_bg,
     }
-    headers = {"Authorization": f"Api-Key {BTEN_API_KEY or os.getenv('API_KEY')}"}
+    if not BTEN_API_KEY:
+        headers = {"Authorization": f"Api-Key {os.getenv('API_KEY')}"}
+    else:
+        headers = {"Authorization": f"Api-Key {BTEN_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()
@@ -61,17 +83,20 @@ def call_baseten_generate(image: Image.Image, prompt: str, steps: int, strength:
         return None
 
 
-# ================== MODE CONFIG =====================
+# Mode defaults for each tab.
 
-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": {
@@ -80,7 +105,12 @@ MODE_DEFAULTS: dict[str, Mode] = {
         "default_strength": 1.2,
         "default_height": 512,
         "default_width": 512,
-        "models": ["zendsd_512_146000", "subject_99000_512", "zen_26000_512"],
+        "models": [
+            "zendsd_512_146000",
+            "subject_99000_512",
+            # "zen_pers_11000",
+            "zen_26000_512",
+        ],
         "remove_bg": True,
     },
     "Background Generation": {
@@ -89,7 +119,20 @@ MODE_DEFAULTS: dict[str, Mode] = {
         "default_strength": 1.2,
         "default_height": 1024,
         "default_width": 1024,
-        "models": ["bgwlight_15000_1024", "bg_canny_58000_1024", "gen_back_7000_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",
+        ],
         "remove_bg": True,
     },
     "Canny": {
@@ -107,7 +150,9 @@ 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": {
@@ -116,87 +161,63 @@ MODE_DEFAULTS: dict[str, Mode] = {
         "default_strength": 1.2,
         "default_height": 1024,
         "default_width": 1024,
-        "models": ["deblurr_1024_10000"],
+        "models": ["deblurr_1024_10000"],  # "slight_deblurr_18000",
         "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"><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>
+    <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>
 </div>
 """
 
+defaults = MODE_DEFAULTS["Subject Generation"]
+
+
 with gr.Blocks(title="🌍 ZenCtrl") as demo:
     gr.HTML(header)
-    gr.Markdown("# ZenCtrl Demo")
-
+    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.
+        """
+    )
     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):
-                        input_image = gr.Image(label="Input Image", type="pil")
+                    with gr.Column(scale=2, min_width=370):
+                        input_image = gr.Image(
+                            label="Upload Image",
+                            type="pil",
+                            scale=3,
+                            height=370,
+                            min_width=100,
+                        )
                         generate_button = gr.Button("Generate")
-                        with gr.Blocks():
+                        with gr.Blocks(title="Options"):
                             model_dropdown = gr.Dropdown(
                                 label="Model",
                                 choices=defaults["models"],
@@ -208,12 +229,20 @@ with gr.Blocks(title="🌍 ZenCtrl") as demo:
                             )
 
                     with gr.Column(scale=2):
-                        output_image = gr.Image(label="Generated Image", type="pil")
+                        output_image = gr.Image(
+                            label="Generated Image",
+                            type="pil",
+                            height=573,
+                            scale=4,
+                            min_width=100,
+                        )
 
+                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(
@@ -242,8 +271,26 @@ 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,
@@ -258,16 +305,9 @@ 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()
+    demo.launch()