Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/example_1.png filter=lfs diff=lfs merge=lfs -text
+assets/example_3.png filter=lfs diff=lfs merge=lfs -text
+assets/example_5.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,42 @@
 ---
 title: Stable Layers
-emoji: 🐨
-colorFrom: pink
-colorTo: indigo
+emoji: 🎨
+colorFrom: yellow
+colorTo: purple
 sdk: gradio
 sdk_version: 6.22.0
-python_version: '3.12'
 app_file: app.py
 pinned: false
+short_description: Decompose images into editable RGBA layers
+python_version: "3.12"
+models:
+  - Qwen/Qwen-Image-Layered
+  - StabilityLabs/Stable-Layers
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Stable Layers — Image Decomposition
+
+Decompose any image into a stack of editable RGBA layers (background + objects) using
+[StabilityLabs/Stable-Layers](https://huggingface.co/StabilityLabs/Stable-Layers), a LoRA
+adapter fine-tuned on the [Qwen/Qwen-Image-Layered](https://huggingface.co/Qwen/Qwen-Image-Layered)
+base pipeline.
+
+## How it works
+
+Upload an image and the model decomposes it into ordered, back-to-front RGBA layers
+suitable for compositing and editing. Layer 0 is the background (inpainted behind the
+removed objects); higher-index layers sit on top.
+
+**Inference settings** (locked to the model card's recommended recipe):
+- Heun sampler (2nd order) — 50 steps
+- CFG 1.0 (disabled)
+- 640px max dimension
+- 4 layers by default
+
+The transparent RGBA toggle controls whether layers are saved with real alpha (for
+compositing in an editor) or composited onto white (for easy eyeballing).
+
+## Credits
+
+Example images sourced from the [Qwen-Image-Layered demo Space](https://huggingface.co/spaces/Qwen/Qwen-Image-Layered)
+(Apache 2.0 license).

app.py

@@ -0,0 +1,473 @@
+import spaces  # MUST be first
+import os
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import tempfile
+import zipfile
+import random
+import numpy as np
+import torch
+import gradio as gr
+from PIL import Image
+from pathlib import Path
+
+# ---------------------------------------------------------------------------
+# Constants — from decompose.py (reference inference script)
+# ---------------------------------------------------------------------------
+BASE_MODEL = "Qwen/Qwen-Image-Layered"
+LORA_REPO = "StabilityLabs/Stable-Layers"
+LORA_SUBFOLDER = "model"  # adapter_config.json + adapter_model.safetensors live under model/
+
+PROMPT_TEMPLATE = (
+    "system\nDescribe the image by detailing the color, shape, size, "
+    "texture, quantity, text, spatial relationships of the objects and background:"
+    "\nuser\n{}\nassistant\n"
+)
+PROMPT_TEMPLATE_DROP_IDX = 34
+DEFAULT_PROMPT = "a clean, well composed image"
+VAE_SCALE_FACTOR = 8
+MAX_SEED = np.iinfo(np.int32).max
+
+# ---------------------------------------------------------------------------
+# Latent helpers — verbatim from decompose.py
+# ---------------------------------------------------------------------------
+
+def rgb_to_rgba(image):
+    b, _, h, w = image.shape
+    alpha = torch.ones(b, 1, h, w, device=image.device, dtype=image.dtype)
+    return torch.cat([image, alpha], dim=1)
+
+
+def normalize_latents(latents, vae):
+    mean = torch.tensor(vae.config.latents_mean).view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
+    std = torch.tensor(vae.config.latents_std).view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
+    return (latents - mean) / std
+
+
+def denormalize_latents(latents, vae):
+    mean = torch.tensor(vae.config.latents_mean).view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
+    inv_std = (1.0 / torch.tensor(vae.config.latents_std)).view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
+    return latents / inv_std + mean
+
+
+def pack_latents(latents, batch_size, num_channels, height, width, num_frames):
+    latents = latents.view(batch_size, num_frames, num_channels, height // 2, 2, width // 2, 2)
+    latents = latents.permute(0, 1, 3, 5, 2, 4, 6)
+    return latents.reshape(batch_size, num_frames * (height // 2) * (width // 2), num_channels * 4)
+
+
+def unpack_latents(latents, height, width, num_layers, vae_scale_factor=VAE_SCALE_FACTOR):
+    batch_size, _, channels = latents.shape
+    frames = num_layers + 1
+    h = 2 * (int(height) // (vae_scale_factor * 2))
+    w = 2 * (int(width) // (vae_scale_factor * 2))
+    latents = latents.view(batch_size, frames, h // 2, w // 2, channels // 4, 2, 2)
+    latents = latents.permute(0, 1, 4, 2, 5, 3, 6)
+    latents = latents.reshape(batch_size, frames, channels // 4, h, w)
+    return latents.permute(0, 2, 1, 3, 4)
+
+
+@torch.no_grad()
+def encode_condition_image(vae, image):
+    """RGB image tensor in [-1,1], (B,3,H,W) -> packed condition latents."""
+    b = image.shape[0]
+    vae_input_ch = getattr(vae.config, "input_channels", 3)
+    if vae_input_ch == 4 and image.shape[1] == 3:
+        image = rgb_to_rgba(image)
+    elif vae_input_ch == 3 and image.shape[1] == 4:
+        image = image[:, :3]
+
+    image_5d = image.unsqueeze(2).to(dtype=vae.dtype)
+    dist = vae.encode(image_5d)
+    if hasattr(dist, "latent_dist"):
+        latents = dist.latent_dist.mode()
+    elif hasattr(dist, "mode"):
+        latents = dist.mode()
+    else:
+        latents = dist
+
+    latents = normalize_latents(latents, vae)
+    z_dim, lh, lw = latents.shape[1], latents.shape[3], latents.shape[4]
+    latents = latents.permute(0, 2, 1, 3, 4)
+    return pack_latents(latents, b, z_dim, lh, lw, 1).to(dtype=torch.bfloat16)
+
+
+@torch.no_grad()
+def decode_layers(vae, latents, height, width, num_layers):
+    b = latents.shape[0]
+    unpacked = unpack_latents(latents, height, width, num_layers)
+    unpacked = denormalize_latents(unpacked, vae)
+    layer_latents = unpacked[:, :, 1:].permute(0, 2, 1, 3, 4)  # drop base frame
+    _, _, c, h, w = layer_latents.shape
+    decoded = vae.decode(layer_latents.reshape(b * num_layers, c, 1, h, w).to(dtype=vae.dtype),
+                         return_dict=False)[0]
+    decoded = decoded.squeeze(2).float().clamp(-1.0, 1.0)
+    _, c_out, h_out, w_out = decoded.shape
+    return decoded.reshape(b, num_layers, c_out, h_out, w_out)
+
+
+def composite_layers(layers):
+    """Porter-Duff 'over', back-to-front."""
+    if layers.shape[2] == 3:
+        return layers.mean(dim=1)
+    result = torch.zeros_like(layers[:, 0, :3])
+    for i in range(layers.shape[1]):
+        rgb = layers[:, i, :3]
+        alpha = (layers[:, i, 3:4] + 1.0) / 2.0
+        result = rgb * alpha + result * (1.0 - alpha)
+    return result
+
+
+def tensor_to_pil(t, mode="RGB"):
+    arr = ((t.clamp(-1, 1) + 1) / 2 * 255).byte().permute(1, 2, 0).cpu().numpy()
+    return Image.fromarray(arr, mode=mode)
+
+
+def layer_to_pil_white_bg(layer):
+    if layer.shape[0] == 4:
+        rgb, alpha = layer[:3], (layer[3:4] + 1.0) / 2.0
+        return tensor_to_pil(rgb * alpha + torch.ones_like(rgb) * (1.0 - alpha))
+    return tensor_to_pil(layer)
+
+
+def layer_to_pil_rgba(layer):
+    layer = layer.clamp(-1, 1).float()
+    rgba = torch.cat([(layer[:3] + 1) / 2, (layer[3:4] + 1) / 2], dim=0)
+    arr = (rgba * 255).byte().permute(1, 2, 0).cpu().numpy()
+    return Image.fromarray(arr, mode="RGBA")
+
+
+# ---------------------------------------------------------------------------
+# Denoising — Heun (2nd order) — verbatim from decompose.py
+# ---------------------------------------------------------------------------
+
+def transformer_forward(transformer, hidden_states, timestep, encoder_hidden_states,
+                        img_shapes, encoder_hidden_states_mask=None, additional_t_cond=None):
+    return transformer(
+        hidden_states=hidden_states,
+        timestep=timestep / 1000,
+        encoder_hidden_states=encoder_hidden_states,
+        encoder_hidden_states_mask=encoder_hidden_states_mask,
+        img_shapes=img_shapes,
+        guidance=None,
+        additional_t_cond=additional_t_cond,
+        return_dict=False,
+    )[0]
+
+
+def build_img_shapes(num_layers, packed_h, packed_w):
+    return [[*[(1, packed_h, packed_w) for _ in range(num_layers + 1)], (1, packed_h, packed_w)]]
+
+
+@torch.no_grad()
+def denoise(transformer, scheduler, latents, timesteps, prompt_embeds, prompt_mask,
+            img_shapes, condition_latents, guidance_scale,
+            neg_embeds=None, neg_mask=None, additional_t_cond=None):
+    """Heun's method (2nd order) over the scheduler's sigmas."""
+    gen_seq_len = latents.shape[1]
+    do_cfg = guidance_scale > 1.0 and neg_embeds is not None
+
+    def velocity(lat, t_val):
+        ts = t_val.expand(1).to(lat.dtype)
+        model_input = lat if condition_latents is None else torch.cat([lat, condition_latents], dim=1)
+        v = transformer_forward(transformer, model_input, ts, prompt_embeds,
+                                img_shapes, prompt_mask, additional_t_cond)[:, :gen_seq_len]
+        if do_cfg:
+            v_neg = transformer_forward(transformer, model_input, ts, neg_embeds,
+                                        img_shapes, neg_mask, additional_t_cond)[:, :gen_seq_len]
+            v = v_neg + guidance_scale * (v - v_neg)
+        return v
+
+    sigmas = scheduler.sigmas
+    for i, t in enumerate(timesteps):
+        sigma = sigmas[i]
+        sigma_next = sigmas[i + 1] if i + 1 < len(sigmas) else torch.tensor(0.0, device=t.device)
+        dt = sigma_next - sigma
+
+        v1 = velocity(latents, t)
+        latents_mid = latents + dt * v1
+        if sigma_next > 0:                       # Heun corrector
+            v2 = velocity(latents_mid, sigma_next * 1000)
+            latents = latents + dt * 0.5 * (v1 + v2)
+        else:
+            latents = latents_mid
+    return latents
+
+
+# ---------------------------------------------------------------------------
+# Resize helper — from decompose.py
+# ---------------------------------------------------------------------------
+
+def compute_aspect_resize(orig_w, orig_h, max_size):
+    """Keep aspect ratio, max dim = max_size, both dims divisible by 16."""
+    scale = max_size / max(orig_w, orig_h)
+    return (max(int(round(orig_w * scale / 16)) * 16, 16),
+            max(int(round(orig_h * scale / 16)) * 16, 16))
+
+
+# ---------------------------------------------------------------------------
+# Model loading at module scope (ZeroGPU rule)
+# ---------------------------------------------------------------------------
+
+print("Loading base model: " + BASE_MODEL)
+from diffusers import DiffusionPipeline
+from peft import PeftModel
+
+pipe = DiffusionPipeline.from_pretrained(
+    BASE_MODEL, torch_dtype=torch.bfloat16,
+    trust_remote_code=True,
+)
+transformer = pipe.transformer.to("cuda").eval()
+vae = pipe.vae.to("cuda").eval()
+scheduler = pipe.scheduler
+tokenizer = getattr(pipe, "tokenizer", None)
+text_encoder = getattr(pipe, "text_encoder", None)
+del pipe
+transformer.requires_grad_(False)
+vae.requires_grad_(False)
+
+print("Loading LoRA: " + LORA_REPO + "/" + LORA_SUBFOLDER)
+transformer = PeftModel.from_pretrained(transformer, LORA_REPO, subfolder=LORA_SUBFOLDER)
+transformer.eval()
+
+in_channels = getattr(transformer.config, "in_channels", 64)
+additional_t_cond = torch.zeros(1, dtype=torch.long, device="cuda")
+
+if text_encoder is None or tokenizer is None:
+    raise RuntimeError("base model has no text encoder/tokenizer")
+text_encoder = text_encoder.to("cuda").eval()
+text_encoder.requires_grad_(False)
+
+import inspect
+supports_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters)
+
+
+@torch.no_grad()
+def encode_prompt(text):
+    tokens = tokenizer([PROMPT_TEMPLATE.format(text)], padding=True,
+                       return_tensors="pt").to(text_encoder.device)
+    out = text_encoder(input_ids=tokens.input_ids,
+                       attention_mask=tokens.attention_mask,
+                       output_hidden_states=True)
+    hidden = out.hidden_states[-1]
+    mask_bool = tokens.attention_mask.bool()
+    lengths = mask_bool.sum(dim=1)
+    chunks = torch.split(hidden[mask_bool], lengths.tolist(), dim=0)
+    chunks = [c[PROMPT_TEMPLATE_DROP_IDX:] for c in chunks]   # drop system prefix
+    max_len = max(c.size(0) for c in chunks)
+    embeds = torch.stack([torch.cat([c, c.new_zeros(max_len - c.size(0), c.size(1))]) for c in chunks])
+    mask = torch.stack([
+        torch.cat([torch.ones(c.size(0), dtype=torch.long, device=c.device),
+                   torch.zeros(max_len - c.size(0), dtype=torch.long, device=c.device)])
+        for c in chunks
+    ])
+    return embeds.to(dtype=torch.bfloat16, device="cuda"), mask.to("cuda")
+
+
+# Pre-encode default prompt and negative prompt at module scope
+prompt_embeds, prompt_mask = encode_prompt(DEFAULT_PROMPT)
+neg_embeds, neg_mask = encode_prompt("")
+
+print("Model loaded successfully.")
+
+
+# ---------------------------------------------------------------------------
+# Inference
+# ---------------------------------------------------------------------------
+
+@spaces.GPU(duration=300)
+def decompose(input_image, num_layers=4, transparent=True, seed=42,
+              randomize_seed=False, prompt="", progress=gr.Progress(track_tqdm=True)):
+    """Decompose an image into editable RGBA layers using Stable Layers.
+
+    Args:
+        input_image: Source image to decompose into layers.
+        num_layers: Number of layers to decompose into (2-6 recommended).
+        transparent: If True, output RGBA layers with real alpha; if False, composite onto white.
+        seed: Random seed for reproducibility.
+        randomize_seed: If True, pick a random seed each run.
+        prompt: Optional text prompt describing the image (default works well).
+    """
+    if input_image is None:
+        raise gr.Error("Please upload an image first.")
+
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    seed = int(seed)
+
+    if isinstance(input_image, np.ndarray):
+        pil_image = Image.fromarray(input_image)
+    else:
+        pil_image = input_image
+    image = pil_image.convert("RGB")
+
+    tw, th = compute_aspect_resize(*image.size, 640)
+    source = image.resize((tw, th), Image.LANCZOS)
+
+    packed_h = (th // VAE_SCALE_FACTOR) // 2
+    packed_w = (tw // VAE_SCALE_FACTOR) // 2
+
+    img_t = torch.from_numpy(np.array(source)).permute(2, 0, 1).float() / 127.5 - 1.0
+    img_t = img_t.unsqueeze(0).to("cuda")
+    condition_latents = encode_condition_image(vae, img_t)
+
+    # dynamic-shift mu from the condition sequence length
+    mu = (condition_latents.shape[1] / (256 * 256 / 16 / 16)) ** 0.5
+    if supports_sigmas:
+        scheduler.set_timesteps(sigmas=np.linspace(1.0, 0, 50 + 1)[:-1],
+                                device="cuda", mu=mu)
+    else:
+        scheduler.set_timesteps(num_inference_steps=50, device="cuda")
+
+    img_shapes = build_img_shapes(num_layers, packed_h, packed_w)
+    seq_len = (num_layers + 1) * packed_h * packed_w
+
+    gen = torch.Generator(device="cuda").manual_seed(seed)
+    latents = torch.randn(1, seq_len, in_channels, device="cuda",
+                          dtype=torch.bfloat16, generator=gen)
+
+    # Use custom prompt if provided, else use pre-encoded default
+    if prompt and prompt.strip():
+        p_embeds, p_mask = encode_prompt(prompt.strip())
+    else:
+        p_embeds, p_mask = prompt_embeds, prompt_mask
+
+    latents = denoise(transformer, scheduler, latents, scheduler.timesteps,
+                      p_embeds, p_mask, img_shapes, condition_latents,
+                      1.0, neg_embeds, neg_mask, additional_t_cond)
+
+    decoded = decode_layers(vae, latents, th, tw, num_layers)
+    composite = composite_layers(decoded)
+
+    # Build output gallery: [composite] + [layer_0..N-1]
+    output_images = []
+    composite_pil = tensor_to_pil(composite[0])
+    output_images.append(("composite", composite_pil))
+
+    for li in range(num_layers):
+        layer = decoded[0, li]
+        pil = layer_to_pil_rgba(layer) if transparent else layer_to_pil_white_bg(layer)
+        label = "background" if li == 0 else f"layer {li}"
+        output_images.append((label, pil))
+
+    # Create a ZIP file for download
+    temp_files = []
+    with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp:
+        with zipfile.ZipFile(tmp.name, 'w', zipfile.ZIP_DEFLATED) as zipf:
+            composite_pil.save(tmp.name.replace('.zip', '_composite.png'))
+            zipf.write(tmp.name.replace('.zip', '_composite.png'), "composite.png")
+            for i, (label, pil) in enumerate(output_images[1:]):
+                fname = f"layer_{i}.png"
+                pil_path = tmp.name.replace('.zip', f'_{fname}')
+                pil.save(pil_path)
+                zipf.write(pil_path, fname)
+                temp_files.append(pil_path)
+            temp_files.append(tmp.name.replace('.zip', '_composite.png'))
+        zip_path = tmp.name
+
+    # Clean up temp files
+    for tf in temp_files:
+        try:
+            os.unlink(tf)
+        except Exception:
+            pass
+
+    # Return gallery as list of (image, caption) tuples + zip file + seed
+    gallery_items = [(pil, label) for label, pil in output_images]
+    return gallery_items, zip_path, seed
+
+
+# ---------------------------------------------------------------------------
+# Gradio UI
+# ---------------------------------------------------------------------------
+
+CSS = """
+#col-container { max-width: 1200px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+with gr.Blocks(title="Stable Layers — Image Decomposition", css=CSS) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(
+            "# 🎨 Stable Layers\n"
+            "Decompose any image into editable RGBA layers using "
+            "[StabilityLabs/Stable-Layers](https://huggingface.co/StabilityLabs/Stable-Layers), "
+            "a LoRA adapter on [Qwen/Qwen-Image-Layered](https://huggingface.co/Qwen/Qwen-Image-Layered)."
+        )
+        gr.Markdown(
+            "Upload an image and the model decomposes it into **background + object layers** "
+            "(back-to-front, Porter-Duff 'over' compositing). "
+            "Each layer is an editable RGBA image you can download as a ZIP."
+        )
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                input_image = gr.Image(
+                    label="Input Image",
+                    type="pil",
+                    image_mode="RGB",
+                )
+                num_layers = gr.Slider(
+                    label="Number of layers",
+                    minimum=2,
+                    maximum=6,
+                    step=1,
+                    value=4,
+                )
+                transparent = gr.Checkbox(
+                    label="Transparent RGBA layers (real alpha for compositing)",
+                    value=True,
+                )
+                run_button = gr.Button("Decompose!", variant="primary", size="lg")
+
+                with gr.Accordion("Advanced Settings", open=False):
+                    prompt = gr.Textbox(
+                        label="Prompt (optional — default works well)",
+                        placeholder="a clean, well composed image",
+                        value="",
+                        lines=2,
+                    )
+                    seed = gr.Slider(
+                        label="Seed",
+                        minimum=0,
+                        maximum=MAX_SEED,
+                        step=1,
+                        value=42,
+                    )
+                    randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
+                    gr.Markdown(
+                        "_Heun sampler · 50 steps · CFG 1.0 · 640px are locked to the "
+                        "recommended settings from the model card._"
+                    )
+
+            with gr.Column(scale=2):
+                gallery = gr.Gallery(
+                    label="Decomposed Layers",
+                    columns=3,
+                    height="auto",
+                    format="png",
+                )
+                zip_file = gr.File(label="Download all layers (ZIP)")
+                used_seed = gr.Number(label="Seed used", precision=0, interactive=False)
+
+        gr.Examples(
+            examples=[
+                ["assets/example_1.png", 4, True],
+                ["assets/example_3.png", 4, True],
+                ["assets/example_5.png", 4, True],
+            ],
+            inputs=[input_image, num_layers, transparent],
+            outputs=[gallery, zip_file, used_seed],
+            fn=decompose,
+            cache_examples=True,
+            cache_mode="lazy",
+        )
+
+    run_button.click(
+        fn=decompose,
+        inputs=[input_image, num_layers, transparent, seed, randomize_seed, prompt],
+        outputs=[gallery, zip_file, used_seed],
+    )
+
+if __name__ == "__main__":
+    demo.launch(theme=gr.themes.Citrus(), mcp_server=True)

requirements.txt

@@ -0,0 +1,8 @@
+git+https://github.com/huggingface/diffusers
+transformers
+accelerate
+peft
+safetensors
+sentencepiece
+torchvision
+numpy