update app

app.py

@@ -83,8 +83,8 @@ class OrangeRedTheme(Soft):
 
 orange_red_theme = OrangeRedTheme()
 
-dtype   = torch.bfloat16
-device  = "cuda" if torch.cuda.is_available() else "cpu"
+dtype  = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 MAX_SEED       = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
@@ -114,78 +114,99 @@ pipe_small_decoder.enable_model_cpu_offload()
 pipe_lock_standard = threading.Lock()
 pipe_lock_small    = threading.Lock()
 
+
+# ── dimension helper ────────────────────────────────────────────────────────
+def calc_dimensions(pil_img: Image.Image):
+    """
+    Given a PIL image return (width, height) snapped to multiples of 8,
+    fitting within 1024 px on the long side, min 256 px on each side.
+    Uses round() so we match the reference app exactly.
+    """
+    iw, ih = pil_img.size
+    aspect = iw / ih
+
+    if aspect >= 1:          # landscape / square
+        new_width  = 1024
+        new_height = int(round(1024 / aspect))
+    else:                    # portrait
+        new_height = 1024
+        new_width  = int(round(1024 * aspect))
+
+    # snap to 8-pixel grid with round(), clamp to [256, 1024]
+    new_width  = max(256, min(1024, round(new_width  / 8) * 8))
+    new_height = max(256, min(1024, round(new_height / 8) * 8))
+    return new_width, new_height
+
+
 def update_dimensions_from_image(image_list):
-    if image_list is None or len(image_list) == 0:
+    """
+    Called by the gallery .upload() event.
+    Returns updated slider values for width and height.
+    """
+    if not image_list:
         return 1024, 1024
 
+    # gallery items arrive as PIL images when type="pil"
     item = image_list[0]
     img  = item[0] if isinstance(item, tuple) else item
 
     if isinstance(img, str):
         img = Image.open(img).convert("RGB")
+    elif not isinstance(img, Image.Image):
+        return 1024, 1024
 
-    iw, ih       = img.size
-    aspect_ratio = iw / ih
+    return calc_dimensions(img)
 
-    if aspect_ratio >= 1:
-        new_width  = 1024
-        new_height = int(1024 / aspect_ratio)
-    else:
-        new_height = 1024
-        new_width  = int(1024 * aspect_ratio)
-
-    new_width  = max(256, min(1024, round(new_width  / 8) * 8))
-    new_height = max(256, min(1024, round(new_height / 8) * 8))
-    return new_width, new_height
-
-def get_example_items():
-    example_prompts = {
-        "1.jpg": "Change the weather to stormy.",
-        "2.jpg": "Transform the scene into a snowy winter day while preserving the original subject identity, framing, and composition.",
-        "3.jpg": "Relight the image with soft golden sunset lighting while keeping all structures and subject details consistent.",
-        "4.jpg": "Make the texture high-resolution.",
-    }
-    items = []
-    if EXAMPLES_DIR.exists():
-        for name in sorted(os.listdir(EXAMPLES_DIR)):
-            if name.lower().endswith((".png", ".jpg", ".jpeg", ".webp")):
-                items.append({
-                    "file":   name,
-                    "path":   str(EXAMPLES_DIR / name),
-                    "prompt": example_prompts.get(
-                        name, "Edit this image while preserving composition."
-                    ),
-                })
-    return items
 
-def parse_input_images(input_images):
-    """Safely parse gallery / filepath / PIL inputs → list[PIL.Image] or None."""
+# ── image parser ─────────────────────────────────────────────────────────────
+def parse_and_resize_images(input_images, width: int, height: int):
+    """
+    Parse the gallery input and resize every frame to (width, height).
+    Returns a list[PIL.Image] or None.
+    """
     if input_images is None:
         return None
+
+    raw_list = []
+
     if isinstance(input_images, str):
-        return [Image.open(input_images).convert("RGB")] if os.path.exists(input_images) else None
-    if isinstance(input_images, list) and len(input_images) > 0:
-        parsed = []
+        if os.path.exists(input_images):
+            raw_list = [Image.open(input_images).convert("RGB")]
+    elif isinstance(input_images, Image.Image):
+        raw_list = [input_images.convert("RGB")]
+    elif isinstance(input_images, list):
         for item in input_images:
             try:
                 src = item[0] if isinstance(item, tuple) else item
                 if isinstance(src, str):
-                    parsed.append(Image.open(src).convert("RGB"))
+                    raw_list.append(Image.open(src).convert("RGB"))
                 elif isinstance(src, Image.Image):
-                    parsed.append(src.convert("RGB"))
+                    raw_list.append(src.convert("RGB"))
                 elif hasattr(src, "name"):
-                    parsed.append(Image.open(src.name).convert("RGB"))
+                    raw_list.append(Image.open(src.name).convert("RGB"))
             except Exception as e:
                 print(f"Skipping invalid image: {e}")
-        return parsed or None
-    return None
 
+    if not raw_list:
+        return None
+
+    # ── KEY FIX: resize every image to the exact pipeline dimensions ──
+    resized = [
+        img.resize((width, height), Image.LANCZOS)
+        for img in raw_list
+    ]
+    return resized
+
+
+# ── pipeline runner ───────────────────────────────────────────────────────────
 def run_pipeline(pipe, lock, kwargs, seed):
     with lock:
         gen    = torch.Generator(device="cpu").manual_seed(seed)
         result = pipe(**kwargs, generator=gen).images[0]
     return result
 
+
+# ── main inference ────────────────────────────────────────────────────────────
 @spaces.GPU(duration=120)
 def infer(
     prompt,
@@ -201,13 +222,38 @@ def infer(
     gc.collect()
     torch.cuda.empty_cache()
 
-    if not prompt or prompt.strip() == "":
+    if not prompt or not prompt.strip():
         raise gr.Error("Please enter a prompt.")
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
-    image_list = parse_input_images(input_images)
+    # ── width / height: derive from the first uploaded image if present ──
+    image_list = None
+    if input_images:
+        # Re-derive dimensions from the actual first image so they are
+        # always consistent with what the pipeline will receive.
+        item = (
+            input_images[0][0]
+            if isinstance(input_images[0], tuple)
+            else input_images[0]
+        )
+        if isinstance(item, str):
+            first_pil = Image.open(item).convert("RGB")
+        elif isinstance(item, Image.Image):
+            first_pil = item.convert("RGB")
+        else:
+            first_pil = None
+
+        if first_pil is not None:
+            width, height = calc_dimensions(first_pil)
+
+        # parse + resize all images to the final (width, height)
+        image_list = parse_and_resize_images(input_images, width, height)
+
+    # ensure dims are multiples of 8 even for text-only runs
+    width  = max(256, min(MAX_IMAGE_SIZE, round(int(width)  / 8) * 8))
+    height = max(256, min(MAX_IMAGE_SIZE, round(int(height) / 8) * 8))
 
     shared_kwargs = dict(
         prompt=prompt,
@@ -222,8 +268,12 @@ def infer(
     progress(0.05, desc="⚡ Launching both pipelines simultaneously...")
 
     with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
-        future_std   = executor.submit(run_pipeline, pipe_standard,     pipe_lock_standard, shared_kwargs, seed)
-        future_small = executor.submit(run_pipeline, pipe_small_decoder, pipe_lock_small,    shared_kwargs, seed)
+        future_std   = executor.submit(
+            run_pipeline, pipe_standard,     pipe_lock_standard, shared_kwargs, seed
+        )
+        future_small = executor.submit(
+            run_pipeline, pipe_small_decoder, pipe_lock_small,    shared_kwargs, seed
+        )
         concurrent.futures.wait(
             [future_std, future_small],
             return_when=concurrent.futures.ALL_COMPLETED,
@@ -254,6 +304,27 @@ def infer_example(prompt):
     )
     return out_std, out_small, seed_used
 
+
+def get_example_items():
+    example_prompts = {
+        "1.jpg": "Change the weather to stormy.",
+        "2.jpg": "Transform the scene into a snowy winter day while preserving the original subject identity, framing, and composition.",
+        "3.jpg": "Relight the image with soft golden sunset lighting while keeping all structures and subject details consistent.",
+        "4.jpg": "Make the texture high-resolution.",
+    }
+    items = []
+    if EXAMPLES_DIR.exists():
+        for name in sorted(os.listdir(EXAMPLES_DIR)):
+            if name.lower().endswith((".png", ".jpg", ".jpeg", ".webp")):
+                items.append({
+                    "file":   name,
+                    "path":   str(EXAMPLES_DIR / name),
+                    "prompt": example_prompts.get(
+                        name, "Edit this image while preserving composition."
+                    ),
+                })
+    return items
+
 EXAMPLE_ITEMS = get_example_items()
 
 css = """
@@ -284,11 +355,13 @@ with gr.Blocks() as demo:
             elem_id="main-title",
         )
         gr.Markdown(
-            "Compare **FLUX.2-klein-4B** side-by-side with [samll decoder](https://huggingface.co/black-forest-labs/FLUX.2-small-decoder)."
+            "Compare **FLUX.2-klein-4B** side-by-side with "
+            "[small decoder](https://huggingface.co/black-forest-labs/FLUX.2-small-decoder)."
         )
 
         with gr.Row(equal_height=True):
 
+            # ── LEFT COLUMN: inputs ─────────────────────────────────────────
             with gr.Column():
                 input_images = gr.Gallery(
                     label="Input Images",
@@ -304,9 +377,10 @@ with gr.Blocks() as demo:
                     show_label=True,
                     placeholder="e.g., A black cat holding a sign that says hello world...",
                 )
-                    
+
                 run_button = gr.Button("Run Comparison", variant="primary")
 
+            # ── RIGHT COLUMN: outputs ───────────────────────────────────────
             with gr.Column():
                 with gr.Row():
                     with gr.Column():
@@ -317,7 +391,6 @@ with gr.Blocks() as demo:
                             format="png",
                             height=250,
                         )
-
                     with gr.Column():
                         result_small = gr.Image(
                             label="Small Decoder",
@@ -329,7 +402,7 @@ with gr.Blocks() as demo:
 
                 seed_output = gr.Number(label="Seed Used", precision=0, visible=False)
 
-                with gr.Accordion("Advanced Settings", open=False, visible=False):
+                with gr.Accordion("Advanced Settings", open=False):
                     seed = gr.Slider(
                         label="Seed",
                         minimum=0,
@@ -390,7 +463,8 @@ with gr.Blocks() as demo:
             "[*](https://huggingface.co/black-forest-labs/FLUX.2-klein-4B) "
             "Experimental Space — FLUX.2 [klein] 4B VAE Decoder Comparison."
         )
-        
+
+    # ── events ────────────────────────────────────────────────────────────────
     input_images.upload(
         fn=update_dimensions_from_image,
         inputs=[input_images],
@@ -415,9 +489,8 @@ with gr.Blocks() as demo:
 
 if __name__ == "__main__":
     demo.queue(max_size=20).launch(
-        theme=orange_red_theme, 
-        mcp_server=True, 
-        css=css,
+        theme=orange_red_theme, css=css,
+        mcp_server=True,
         ssr_mode=False,
         show_error=True,
     )