Update app.py

app.py

@@ -1,31 +1,67 @@
 import os
 import uuid
-import numpy as np
 import random
 import tempfile
 import zipfile
+import numpy as np
 
 import spaces
 import torch
 import gradio as gr
 
 from PIL import Image
-from diffusers import QwenImageLayeredPipeline
 from pptx import Presentation
+from diffusers import QwenImageLayeredPipeline
 
 LOG_DIR = "/tmp/local"
 MAX_SEED = np.iinfo(np.int32).max
 
-# Optional HF login (works in Spaces if you set HF token as secret env var "hf")
+# Reduce allocator fragmentation (new name; old PYTORCH_CUDA_ALLOC_CONF is deprecated)
+os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True")
+
+# Optional HF login (Spaces secret env var "hf")
 from huggingface_hub import login
 login(token=os.environ.get("hf"))
 
-dtype = torch.bfloat16
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# ----------------------------
+# Device / dtype (memory-safe)
+# ----------------------------
+has_cuda = torch.cuda.is_available()
+device = "cuda" if has_cuda else ("mps" if torch.backends.mps.is_available() else "cpu")
+
+# fp16 is typically best for VRAM; CPU uses fp32
+torch_dtype = torch.float16 if device in ("cuda", "mps") else torch.float32
 
+# ----------------------------
+# Load pipeline (avoid CPU RAM spikes)
+# ----------------------------
 pipeline = QwenImageLayeredPipeline.from_pretrained(
-    "Qwen/Qwen-Image-Layered", torch_dtype=dtype
-).to(device)
+    "Qwen/Qwen-Image-Layered",
+    torch_dtype=torch_dtype,
+    low_cpu_mem_usage=True,
+)
+
+# Memory helpers (guarded)
+if hasattr(pipeline, "enable_attention_slicing"):
+    pipeline.enable_attention_slicing()
+
+# This pipeline may NOT expose enable_vae_slicing(), so guard both ways
+if hasattr(pipeline, "enable_vae_slicing"):
+    pipeline.enable_vae_slicing()
+elif hasattr(pipeline, "vae") and hasattr(pipeline.vae, "enable_slicing"):
+    pipeline.vae.enable_slicing()
+
+if device == "cuda":
+    # Best for Spaces: keep CPU RAM lower and avoid huge peak VRAM at startup
+    # (requires accelerate, usually present in Spaces)
+    try:
+        pipeline.enable_model_cpu_offload()
+    except Exception:
+        pipeline.to("cuda")
+elif device == "mps":
+    pipeline.to("mps")
+else:
+    pipeline.to("cpu")
 
 
 def ensure_dirname(path: str):
@@ -33,26 +69,21 @@ def ensure_dirname(path: str):
         os.makedirs(path, exist_ok=True)
 
 
-def random_str(length=8):
-    return uuid.uuid4().hex[:length]
-
-
 def imagelist_to_pptx(img_files):
     with Image.open(img_files[0]) as img:
         img_width_px, img_height_px = img.size
 
     def px_to_emu(px, dpi=96):
         inch = px / dpi
-        emu = inch * 914400
-        return int(emu)
+        return int(inch * 914400)
 
     prs = Presentation()
     prs.slide_width = px_to_emu(img_width_px)
     prs.slide_height = px_to_emu(img_height_px)
 
     slide = prs.slides.add_slide(prs.slide_layouts[6])
-
     left = top = 0
+
     for img_path in img_files:
         slide.shapes.add_picture(
             img_path,
@@ -75,171 +106,214 @@ def _clamp_int(x, default: int, lo: int, hi: int) -> int:
     return max(lo, min(hi, v))
 
 
-# Dynamic duration callable: must accept the same args as infer(). It returns seconds.
+def _safe_open_rgba(img_like):
+    if isinstance(img_like, list):
+        img_like = img_like[0]
+    if isinstance(img_like, str):
+        return Image.open(img_like).convert("RGB").convert("RGBA")
+    if isinstance(img_like, Image.Image):
+        return img_like.convert("RGB").convert("RGBA")
+    if isinstance(img_like, np.ndarray):
+        return Image.fromarray(img_like).convert("RGB").convert("RGBA")
+    raise ValueError(f"Unsupported input_image type: {type(img_like)}")
+
+
+def _update_refine_index_ui(n_layers: int, current_idx: int | None = None):
+    n_layers = max(1, int(n_layers))
+    if current_idx is None:
+        current_idx = 1
+    current_idx = max(1, min(int(current_idx), n_layers))
+    return gr.update(minimum=1, maximum=n_layers, value=current_idx)
+
+
+# Dynamic duration callable: must accept same args as decompose() and refine()
 def get_duration(
     input_image,
-    seed=777,
-    randomize_seed=False,
-    prompt=None,
+    seed=0,
+    randomize_seed=True,
+    prompt="",
     neg_prompt=" ",
     true_guidance_scale=4.0,
     num_inference_steps=50,
-    layer=4,
+    layer=7,
     cfg_norm=True,
     use_en_prompt=True,
-    resolution=640,
-    gpu_duration=1000,
-    refine_enabled=False,
+    resolution=1024,
+    gpu_duration="1000",
     refine_layer_index=1,
     refine_sub_layers=3,
 ):
     return _clamp_int(gpu_duration, default=1000, lo=20, hi=1500)
 
 
-def _normalize_input_image(input_image):
-    if isinstance(input_image, list):
-        input_image = input_image[0]
+@spaces.GPU(duration=get_duration)
+def decompose(
+    input_image,
+    seed=0,
+    randomize_seed=True,
+    prompt="",
+    neg_prompt=" ",
+    true_guidance_scale=4.0,
+    num_inference_steps=50,
+    layer=7,
+    cfg_norm=True,
+    use_en_prompt=True,
+    resolution=1024,
+    gpu_duration="1000",
+    refine_layer_index=1,   # passed in (so we can "clamp" it красиво)
+    refine_sub_layers=3,    # unused here, but kept for duration signature parity
+):
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
 
-    if isinstance(input_image, str):
-        return Image.open(input_image).convert("RGB").convert("RGBA")
-    if isinstance(input_image, Image.Image):
-        return input_image.convert("RGB").convert("RGBA")
-    if isinstance(input_image, np.ndarray):
-        return Image.fromarray(input_image).convert("RGB").convert("RGBA")
+    resolution = _clamp_int(resolution, default=1024, lo=640, hi=1024)
+    if resolution not in (640, 1024):
+        resolution = 1024
 
-    raise ValueError(f"Unsupported input_image type: {type(input_image)}")
+    pil_image = _safe_open_rgba(input_image)
 
+    # Generator on CPU works well with CPU offload too
+    gen = torch.Generator(device="cpu").manual_seed(seed)
 
-def _export_images_to_pptx_and_zip(pil_images, zip_prefix="layer"):
-    temp_files = []
-    for img in pil_images:
+    inputs = {
+        "image": pil_image,
+        "generator": gen,
+        "true_cfg_scale": float(true_guidance_scale),
+        "prompt": prompt if prompt else None,
+        "negative_prompt": neg_prompt,
+        "num_inference_steps": int(num_inference_steps),
+        "num_images_per_prompt": 1,
+        "layers": int(layer),
+        "resolution": int(resolution),
+        "cfg_normalize": bool(cfg_norm),
+        "use_en_prompt": bool(use_en_prompt),
+    }
+
+    print("DECOMPOSE INPUTS:", {k: v for k, v in inputs.items() if k != "image"})
+    print("REQUESTED GPU DURATION:", gpu_duration)
+
+    with torch.inference_mode():
+        out = pipeline(**inputs)
+        output_images = out.images[0]  # list[PIL.Image]
+
+    # Save layers for exports + for refine stage
+    layer_paths = []
+    gallery_out = []
+
+    for img in output_images:
+        gallery_out.append(img)
         tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
         img.save(tmp.name)
-        temp_files.append(tmp.name)
+        layer_paths.append(tmp.name)
 
-    pptx_path = imagelist_to_pptx(temp_files)
+    pptx_path = imagelist_to_pptx(layer_paths)
 
     with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmpzip:
         with zipfile.ZipFile(tmpzip.name, "w", zipfile.ZIP_DEFLATED) as zipf:
-            for i, img_path in enumerate(temp_files):
-                zipf.write(img_path, f"{zip_prefix}_{i+1}.png")
+            for i, p in enumerate(layer_paths):
+                zipf.write(p, f"layer_{i+1}.png")
         zip_path = tmpzip.name
 
-    return pptx_path, zip_path
+    # Reset refined outputs on new decompose
+    refined_gallery = []
+    refined_pptx = None
+    refined_zip = None
+
+    # "совсем красиво": clamp current refine index to new [1..N]
+    refine_index_update = _update_refine_index_ui(len(layer_paths), refine_layer_index)
+
+    return (
+        gallery_out,
+        pptx_path,
+        zip_path,
+        layer_paths,            # gr.State
+        refined_gallery,
+        refined_pptx,
+        refined_zip,
+        refine_index_update,    # update refine slider bounds/value
+    )
 
 
 @spaces.GPU(duration=get_duration)
-def infer(
-    input_image,
-    seed=777,
-    randomize_seed=False,
-    prompt=None,
+def refine_selected_layer(
+    layer_paths,
+    refine_layer_index=1,
+    refine_sub_layers=3,
+    seed=0,
+    randomize_seed=True,
+    prompt="",
     neg_prompt=" ",
     true_guidance_scale=4.0,
     num_inference_steps=50,
-    layer=4,
     cfg_norm=True,
     use_en_prompt=True,
-    resolution=640,
-    gpu_duration=1000,
-    refine_enabled=False,
-    refine_layer_index=1,   # 1-based for UI convenience
-    refine_sub_layers=3,
+    resolution=1024,
+    gpu_duration="1000",
 ):
-    # Seed
+    if not layer_paths:
+        return [], None, None
+
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
-    # Normalize resolution input
-    resolution = _clamp_int(resolution, default=640, lo=640, hi=1024)
+    # Clamp index into existing layers
+    n = len(layer_paths)
+    idx = _clamp_int(refine_layer_index, default=1, lo=1, hi=n) - 1
+
+    sub_layers = _clamp_int(refine_sub_layers, default=3, lo=2, hi=10)
+
+    resolution = _clamp_int(resolution, default=1024, lo=640, hi=1024)
     if resolution not in (640, 1024):
-        resolution = 640
+        resolution = 1024
 
-    # Normalize image input
-    pil_image = _normalize_input_image(input_image)
+    selected_path = layer_paths[idx]
+    selected_layer_img = Image.open(selected_path).convert("RGBA")
 
-    gen_device = "cuda" if torch.cuda.is_available() else "cpu"
-    generator = torch.Generator(device=gen_device).manual_seed(seed)
+    gen = torch.Generator(device="cpu").manual_seed(seed)
 
-    # First pass inputs
     inputs = {
-        "image": pil_image,
-        "generator": generator,
-        "true_cfg_scale": true_guidance_scale,
-        "prompt": prompt,
+        "image": selected_layer_img,
+        "generator": gen,
+        "true_cfg_scale": float(true_guidance_scale),
+        "prompt": prompt if prompt else None,
         "negative_prompt": neg_prompt,
-        "num_inference_steps": num_inference_steps,
+        "num_inference_steps": int(num_inference_steps),
         "num_images_per_prompt": 1,
-        "layers": layer,
-        "resolution": resolution,
-        "cfg_normalize": cfg_norm,
-        "use_en_prompt": use_en_prompt,
+        "layers": int(sub_layers),      # <-- ключевой параметр рекурсивной декомпозиции
+        "resolution": int(resolution),  # тот же resolution (без отдельных опций для refine)
+        "cfg_normalize": bool(cfg_norm),
+        "use_en_prompt": bool(use_en_prompt),
     }
 
-    print("INFER INPUTS:", inputs)
+    print("REFINE INPUTS:", {k: v for k, v in inputs.items() if k != "image"})
     print("REQUESTED GPU DURATION:", gpu_duration)
+    print(f"REFINE: base layer index={idx+1}/{n}, sub_layers={sub_layers}")
 
     with torch.inference_mode():
         out = pipeline(**inputs)
-        output_layers = out.images[0]  # list[PIL.Image]
+        refined_images = out.images[0]
 
-    # Export first pass
-    pptx_path, zip_path = _export_images_to_pptx_and_zip(output_layers, zip_prefix="layer")
-
-    # Optional: Recursive (refine one layer into sub-layers) — no separate steps/resolution/cfg
+    refined_paths = []
     refined_gallery = []
-    refined_pptx = None
-    refined_zip = None
-
-    if refine_enabled and len(output_layers) > 0:
-        idx0 = _clamp_int(refine_layer_index, default=1, lo=1, hi=len(output_layers)) - 1
-        refine_sub_layers = _clamp_int(refine_sub_layers, default=3, lo=2, hi=10)
-
-        selected_layer = output_layers[idx0].convert("RGBA")
-
-        refined_inputs = dict(inputs)  # reuse same params
-        refined_inputs["image"] = selected_layer
-        refined_inputs["layers"] = refine_sub_layers
-
-        print("REFINE ENABLED:", True)
-        print("REFINE LAYER INDEX (1-based):", idx0 + 1)
-        print("REFINE SUB-LAYERS:", refine_sub_layers)
-        print("REFINED INPUTS:", {k: v for k, v in refined_inputs.items() if k != "image"})
+    for img in refined_images:
+        refined_gallery.append(img)
+        tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
+        img.save(tmp.name)
+        refined_paths.append(tmp.name)
 
-        with torch.inference_mode():
-            refined_out = pipeline(**refined_inputs)
-            sub_layers = refined_out.images[0]
+    refined_pptx = imagelist_to_pptx(refined_paths)
 
-        refined_gallery = sub_layers
-        refined_pptx, refined_zip = _export_images_to_pptx_and_zip(sub_layers, zip_prefix=f"sub_layer_{idx0+1}")
+    with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmpzip:
+        with zipfile.ZipFile(tmpzip.name, "w", zipfile.ZIP_DEFLATED) as zipf:
+            for i, p in enumerate(refined_paths):
+                zipf.write(p, f"sub_layer_{i+1}.png")
+        refined_zip = tmpzip.name
 
-    return (
-        output_layers,
-        pptx_path,
-        zip_path,
-        refined_gallery,
-        refined_pptx,
-        refined_zip,
-    )
+    return refined_gallery, refined_pptx, refined_zip
 
 
 ensure_dirname(LOG_DIR)
-
-examples = [
-    "assets/test_images/1.png",
-    "assets/test_images/2.png",
-    "assets/test_images/3.png",
-    "assets/test_images/4.png",
-    "assets/test_images/5.png",
-    "assets/test_images/6.png",
-    "assets/test_images/7.png",
-    "assets/test_images/8.png",
-    "assets/test_images/9.png",
-    "assets/test_images/10.png",
-    "assets/test_images/11.png",
-    "assets/test_images/12.png",
-    "assets/test_images/13.png",
-]
+examples = [f"assets/test_images/{i}.png" for i in range(1, 14)]
 
 with gr.Blocks() as demo:
     with gr.Column(elem_id="col-container"):
@@ -249,10 +323,14 @@ with gr.Blocks() as demo:
         )
         gr.Markdown(
             """
-The text prompt is intended to describe the overall content of the input image—including elements that may be partially occluded (e.g., you may specify the text hidden behind a foreground object). It is not designed to control the semantic content of individual layers explicitly.
+The text prompt is intended to describe the overall content of the input image—including elements that may be partially occluded.
+It is not designed to control the semantic content of individual layers explicitly.
             """
         )
 
+        # State to store layer PNG paths from last Decompose
+        layer_paths_state = gr.State([])
+
         with gr.Row():
             with gr.Column(scale=1):
                 input_image = gr.Image(label="Input Image", image_mode="RGBA")
@@ -260,7 +338,7 @@ The text prompt is intended to describe the overall content of the input image
                 with gr.Accordion("Advanced Settings", open=False):
                     prompt = gr.Textbox(
                         label="Prompt (Optional)",
-                        placeholder="Please enter the prompt to descibe the image. (Optional)",
+                        placeholder="Please enter the prompt to describe the image (optional)",
                         value="",
                         lines=2,
                     )
@@ -271,48 +349,27 @@ The text prompt is intended to describe the overall content of the input image
                         lines=2,
                     )
 
-                    seed = gr.Slider(
-                        label="Seed",
-                        minimum=0,
-                        maximum=MAX_SEED,
-                        step=1,
-                        value=0,
-                    )
+                    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                     randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
 
                     true_guidance_scale = gr.Slider(
-                        label="True guidance scale",
-                        minimum=1.0,
-                        maximum=10.0,
-                        step=0.1,
-                        value=4.0,
+                        label="True guidance scale", minimum=1.0, maximum=10.0, step=0.1, value=4.0
                     )
 
                     num_inference_steps = gr.Slider(
-                        label="Number of inference steps",
-                        minimum=1,
-                        maximum=100,
-                        step=1,
-                        value=50,
+                        label="Number of inference steps", minimum=1, maximum=100, step=1, value=50
                     )
 
-                    layer = gr.Slider(
-                        label="Layers",
-                        minimum=2,
-                        maximum=10,
-                        step=1,
-                        value=7,
-                    )
+                    layer = gr.Slider(label="Layers", minimum=2, maximum=10, step=1, value=7)
 
+                    # default 1024 as you asked earlier
                     resolution = gr.Radio(
                         label="Processing resolution",
                         choices=[640, 1024],
-                        value=640,
+                        value=1024,
                     )
 
-                    cfg_norm = gr.Checkbox(
-                        label="Whether enable CFG normalization", value=True
-                    )
+                    cfg_norm = gr.Checkbox(label="Whether enable CFG normalization", value=True)
                     use_en_prompt = gr.Checkbox(
                         label="Automatic caption language if no prompt provided, True for EN, False for ZH",
                         value=True,
@@ -325,27 +382,24 @@ The text prompt is intended to describe the overall content of the input image
                         placeholder="e.g. 60, 120, 300, 1000, 1500",
                     )
 
-                    gr.Markdown("### Advanced: Recursive decomposition")
-                    refine_enabled = gr.Checkbox(
-                        label="Refine one layer into sub-layers",
-                        value=False,
-                    )
+                decompose_btn = gr.Button("Decompose!", variant="primary")
+
+                with gr.Accordion("Refine layer (Recursive Decomposition)", open=False):
                     refine_layer_index = gr.Slider(
-                        label="Refine layer index (1-based)",
+                        label="Refine layer index (1 = first layer)",
                         minimum=1,
-                        maximum=10,
+                        maximum=7,
                         step=1,
                         value=1,
                     )
                     refine_sub_layers = gr.Slider(
-                        label="Sub-layers (for refined layer)",
+                        label="Sub-layers (how many to split selected layer into)",
                         minimum=2,
                         maximum=10,
                         step=1,
                         value=3,
                     )
-
-                run_button = gr.Button("Decompose!", variant="primary")
+                    refine_btn = gr.Button("Refine selected layer", variant="secondary")
 
             with gr.Column(scale=2):
                 gallery = gr.Gallery(label="Layers", columns=4, rows=1, format="png")
@@ -353,12 +407,13 @@ The text prompt is intended to describe the overall content of the input image
                     export_file = gr.File(label="Download PPTX")
                     export_zip_file = gr.File(label="Download ZIP")
 
-                gr.Markdown("### Refined sub-layers")
+                gr.Markdown("### Refined (sub-layers)")
                 refined_gallery = gr.Gallery(label="Sub-layers", columns=4, rows=1, format="png")
                 with gr.Row():
                     refined_export_file = gr.File(label="Download refined PPTX")
                     refined_export_zip_file = gr.File(label="Download refined ZIP")
 
+    # Examples run Decompose
     gr.Examples(
         examples=examples,
         inputs=[input_image],
@@ -366,18 +421,21 @@ The text prompt is intended to describe the overall content of the input image
             gallery,
             export_file,
             export_zip_file,
+            layer_paths_state,
             refined_gallery,
             refined_export_file,
             refined_export_zip_file,
+            refine_layer_index,  # update slider bounds/value
         ],
-        fn=infer,
+        fn=decompose,
         examples_per_page=14,
         cache_examples=False,
         run_on_click=True,
     )
 
-    run_button.click(
-        fn=infer,
+    # Decompose button
+    decompose_btn.click(
+        fn=decompose,
         inputs=[
             input_image,
             seed,
@@ -391,18 +449,40 @@ The text prompt is intended to describe the overall content of the input image
             use_en_prompt,
             resolution,
             gpu_duration,
-            refine_enabled,
-            refine_layer_index,
-            refine_sub_layers,
+            refine_layer_index,  # so we can clamp nicely after new decomposition
+            refine_sub_layers,   # for duration signature parity
         ],
         outputs=[
             gallery,
             export_file,
             export_zip_file,
+            layer_paths_state,
             refined_gallery,
             refined_export_file,
             refined_export_zip_file,
+            refine_layer_index,  # update slider bounds/value
+        ],
+    )
+
+    # Refine button
+    refine_btn.click(
+        fn=refine_selected_layer,
+        inputs=[
+            layer_paths_state,
+            refine_layer_index,
+            refine_sub_layers,
+            seed,
+            randomize_seed,
+            prompt,
+            neg_prompt,
+            true_guidance_scale,
+            num_inference_steps,
+            cfg_norm,
+            use_en_prompt,
+            resolution,
+            gpu_duration,
         ],
+        outputs=[refined_gallery, refined_export_file, refined_export_zip_file],
     )
 
 if __name__ == "__main__":