Extract transparent assets with BiRefNet

README.md

@@ -67,7 +67,7 @@ If a requested asset has no deterministic hook, images are still generated for r
 
 ## Contract-driven generation and validation
 
-Filenames and IDs do not decide whether an image is a scene or sprite. Dimensions, transparency, composition, camera, quantity, subject count, and silhouette requirements come from the manifest. Transparent single-subject outputs must first arrive from the model on a plain light studio field; scene-like model frames are rejected before alpha extraction. Passing outputs then receive alpha extraction, occupancy normalization, subject-count checks, and a retained-scene check that rejects environmental foreground spread around an isolated subject. Opaque outputs are checked for unwanted transparency. Seamless requests receive an opposite-edge similarity check. Other compositions are not incorrectly rejected for containing multiple subjects.
+Filenames and IDs do not decide whether an image is a scene or sprite. Dimensions, transparency, composition, camera, quantity, subject count, and silhouette requirements come from the manifest. Transparent single-subject outputs receive neural alpha extraction from the pinned MIT-licensed [`ZhengPeng7/BiRefNet`](https://huggingface.co/ZhengPeng7/BiRefNet) model, occupancy normalization, subject-count checks, and a retained-scene check that rejects environmental foreground spread around an isolated subject. Opaque outputs are checked for unwanted transparency. Seamless requests receive an opposite-edge similarity check. Other compositions are not incorrectly rejected for containing multiple subjects.
 
 The image model can still violate these instructions, so every result requires visual review. Single-subject outputs that fail automated checks are regenerated with a fresh seed rather than cropped into a guessed subject. The review UI can regenerate one output or a complete multi-output asset set.
 
@@ -81,15 +81,16 @@ If the Space has an `HF_TOKEN` secret and `USE_HF_PROMPT_PROVIDER=1`, it first u
 
 ## Default neural image pipeline
 
-The deployed Space is designed for Hugging Face ZeroGPU and uses a real diffusion model for both sprites and backgrounds:
+The deployed Space is designed for Hugging Face ZeroGPU and uses real neural models for generation and transparent cutout extraction:
 
 ```text
 segmind/SSD-1B + latent-consistency/lcm-lora-ssd-1b
+ZhengPeng7/BiRefNet (transparent-cutout alpha only)
 ```
 
 SSD-1B is an Apache-2.0 distilled SDXL model. The LCM adapter enables low-step generation. The weights run inside the Space's ZeroGPU allocation, so the app does not consume Hugging Face Inference Provider credits and does not require a paid image API.
 
-Sprites and backgrounds are generated directly from their written prompts. No procedurally drawn sprite, palette, silhouette, camera guide, or background layout is passed into the production model. Sprite prompts use a portrait canvas and positive single-subject composition language; the negative prompt separately rejects sheets, lineups, repeated subjects, multiple views, and surrounding environments. Before alpha extraction, cutout attempts whose boundary is not predominantly light and neutral are rejected because they did not honor the requested studio field. After transparency extraction, sprites with zero or multiple significant foreground components, retained scene-like alpha, implausible silhouettes, or inadequate scale are regenerated with a new seed up to `PRIMARY_SPRITE_ATTEMPTS` times. If every attempt fails, the deployed Space returns an explicit model-generation error rather than returning a contact sheet or environment-contaminated cutout. Background negative prompts reject characters and generic particle overlays.
+Sprites and backgrounds are generated directly from their written prompts. No procedurally drawn sprite, palette, silhouette, camera guide, or background layout is passed into the production image generator. Sprite prompts use a portrait canvas and positive single-subject composition language; the negative prompt separately rejects sheets, lineups, repeated subjects, multiple views, and surrounding environments. BiRefNet predicts only the alpha mask for transparent cutouts while preserving SSD-1B's generated RGB artwork. After neural transparency extraction, sprites with zero or multiple significant foreground components, retained scene-like alpha, implausible silhouettes, or inadequate scale are regenerated with a new seed up to `PRIMARY_SPRITE_ATTEMPTS` times. If every attempt fails, the deployed Space returns an explicit model-generation error rather than returning a contact sheet or environment-contaminated cutout. Background negative prompts reject characters and generic particle overlays.
 
 SSD-1B's CLIP text encoder has a 77-token context window. Production prompts therefore use a conservative word budget, place the camera and single-subject contract first, and compact the free-form description. This prevents those critical constraints from being silently truncated behind verbose style text.
 
@@ -106,6 +107,8 @@ USE_PRIMARY_IMAGE_MODEL=1
 REQUIRE_PRIMARY_IMAGE_MODEL=1
 PRIMARY_IMAGE_MODEL=segmind/SSD-1B
 PRIMARY_LORA_MODEL=latent-consistency/lcm-lora-ssd-1b
+FOREGROUND_MODEL_ID=ZhengPeng7/BiRefNet
+FOREGROUND_MODEL_REVISION=e2bf8e4460fc8fa32bba5ea4d94b3233d367b0e4
 PRIMARY_IMAGE_STEPS=4
 PRIMARY_SPRITE_STEPS=6
 PRIMARY_SPRITE_ATTEMPTS=3

app.py

@@ -167,6 +167,11 @@ class IntegrationReport:
 HF_TOKEN = os.environ.get("HF_TOKEN", "")
 PRIMARY_IMAGE_MODEL = os.environ.get("PRIMARY_IMAGE_MODEL", "segmind/SSD-1B")
 PRIMARY_LORA_MODEL = os.environ.get("PRIMARY_LORA_MODEL", "latent-consistency/lcm-lora-ssd-1b")
+FOREGROUND_MODEL_ID = os.environ.get("FOREGROUND_MODEL_ID", "ZhengPeng7/BiRefNet")
+FOREGROUND_MODEL_REVISION = os.environ.get(
+    "FOREGROUND_MODEL_REVISION",
+    "e2bf8e4460fc8fa32bba5ea4d94b3233d367b0e4",
+)
 PRIMARY_IMAGE_STEPS = int(os.environ.get("PRIMARY_IMAGE_STEPS", "4"))
 PRIMARY_SPRITE_STEPS = int(os.environ.get("PRIMARY_SPRITE_STEPS", "6"))
 PRIMARY_SPRITE_ATTEMPTS = max(1, int(os.environ.get("PRIMARY_SPRITE_ATTEMPTS", "3")))
@@ -198,6 +203,8 @@ CONTROLNET_PIPE = None
 CONTROLNET_ERROR = None
 PRIMARY_TEXT_PIPE = None
 PRIMARY_MODEL_ERROR = None
+FOREGROUND_MODEL = None
+FOREGROUND_MODEL_ERROR = None
 
 
 def gpu_task(duration):
@@ -1949,6 +1956,56 @@ def initialize_primary_image_model() -> None:
         PRIMARY_MODEL_ERROR = short_error(exc)
 
 
+def initialize_foreground_model() -> None:
+    """Load the pinned neural foreground extractor used by transparent cutouts."""
+    global FOREGROUND_MODEL, FOREGROUND_MODEL_ERROR
+    if FOREGROUND_MODEL is not None or FOREGROUND_MODEL_ERROR:
+        return
+    if not USE_PRIMARY_IMAGE_MODEL or not os.environ.get("SPACE_ID"):
+        return
+
+    try:
+        import torch
+        from transformers import AutoModelForImageSegmentation
+
+        model = AutoModelForImageSegmentation.from_pretrained(
+            FOREGROUND_MODEL_ID,
+            revision=FOREGROUND_MODEL_REVISION,
+            trust_remote_code=True,
+            torch_dtype=torch.float16,
+        )
+        FOREGROUND_MODEL = model.eval().to("cuda")
+    except Exception as exc:
+        FOREGROUND_MODEL_ERROR = short_error(exc)
+
+
+def neural_foreground_cutout(image: Image.Image) -> Image.Image:
+    """Use BiRefNet to derive alpha while preserving the generator's RGB artwork."""
+    if FOREGROUND_MODEL is None:
+        raise RuntimeError("neural foreground model is not initialized")
+
+    import numpy as np
+    import torch
+
+    source = image.convert("RGB")
+    sample = source.resize((1024, 1024), Image.LANCZOS)
+    array = np.asarray(sample, dtype=np.float32) / 255.0
+    tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0)
+    mean = torch.tensor((0.485, 0.456, 0.406)).view(1, 3, 1, 1)
+    std = torch.tensor((0.229, 0.224, 0.225)).view(1, 3, 1, 1)
+    tensor = ((tensor - mean) / std).to(device="cuda", dtype=torch.float16)
+
+    with torch.no_grad():
+        prediction = FOREGROUND_MODEL(tensor)[-1].sigmoid()[0].squeeze().float().cpu().numpy()
+    mask = Image.fromarray(np.clip(prediction * 255.0, 0, 255).astype(np.uint8)).resize(
+        source.size,
+        Image.LANCZOS,
+    )
+    cutout = source.convert("RGBA")
+    cutout.putalpha(mask)
+    return cutout
+
+
 def primary_diffusion_prompt(spec: AssetSpec) -> str:
     camera = primary_camera_phrase(spec.camera, is_background_spec(spec))
     if spec.camera == "auto" and spec.camera_instruction and spec.camera_instruction.lower() != "auto":
@@ -1996,10 +2053,18 @@ def primary_diffusion_png(spec: AssetSpec, index: int, run_id: int) -> tuple[byt
     if PRIMARY_TEXT_PIPE is None:
         return None, "primary ZeroGPU diffusion model is not initialized on this runtime"
 
+    is_background = is_background_spec(spec)
+    if not is_background:
+        if FOREGROUND_MODEL is None and not FOREGROUND_MODEL_ERROR:
+            initialize_foreground_model()
+        if FOREGROUND_MODEL_ERROR:
+            return None, f"neural foreground model failed: {FOREGROUND_MODEL_ERROR}"
+        if FOREGROUND_MODEL is None and REQUIRE_PRIMARY_IMAGE_MODEL:
+            return None, "neural foreground model is not initialized on this runtime"
+
     try:
         import torch
 
-        is_background = is_background_spec(spec)
         attempts = 1 if is_background else PRIMARY_SPRITE_ATTEMPTS
         width, height = primary_generation_dimensions(spec)
         last_failure_detail = "the last output did not contain one valid foreground subject"
@@ -2016,12 +2081,18 @@ def primary_diffusion_png(spec: AssetSpec, index: int, run_id: int) -> tuple[byt
                 width=width,
                 height=height,
             ).images[0]
-            if not is_background and has_invalid_cutout_studio_field(image):
-                last_failure_detail = "the last output did not use the required plain light studio field"
-                continue
-            content = polish_diffusion_asset(image, spec)
             if is_background:
+                content = polish_diffusion_asset(image, spec)
                 return content, None
+            if FOREGROUND_MODEL is not None:
+                content = polish_neural_cutout(neural_foreground_cutout(image), spec)
+            else:
+                # Local-development compatibility only. The public fail-closed
+                # runtime requires BiRefNet and cannot enter this branch.
+                if has_invalid_cutout_studio_field(image):
+                    last_failure_detail = "the last output did not use the required plain light studio field"
+                    continue
+                content = polish_diffusion_asset(image, spec)
 
             subject_count = estimated_foreground_subject_count(content)
             if subject_count != 1:
@@ -2050,6 +2121,7 @@ def primary_diffusion_png(spec: AssetSpec, index: int, run_id: int) -> tuple[byt
 # here keeps the weights resident between decorated requests instead of paying the
 # model-load cost on every generation.
 initialize_primary_image_model()
+initialize_foreground_model()
 
 
 def controlnet_background_png(spec: AssetSpec, index: int, run_id: int) -> tuple[bytes | None, str | None]:
@@ -2148,6 +2220,20 @@ def polish_diffusion_asset(image: Image.Image, spec: AssetSpec) -> bytes:
     return out.getvalue()
 
 
+def polish_neural_cutout(image: Image.Image, spec: AssetSpec) -> bytes:
+    """Resize and normalize a neural alpha cutout without changing its generated RGB art."""
+    contained = image.convert("RGBA")
+    contained.thumbnail((spec.width, spec.height), Image.LANCZOS)
+    canvas = Image.new("RGBA", (spec.width, spec.height), (0, 0, 0, 0))
+    offset = ((spec.width - contained.width) // 2, (spec.height - contained.height) // 2)
+    canvas.alpha_composite(contained, dest=offset)
+    if spec.composition in {"single_subject", "icon", "animation_frame"}:
+        canvas = normalize_sprite_foreground(canvas, spec)
+    out = io.BytesIO()
+    canvas.save(out, format="PNG")
+    return out.getvalue()
+
+
 def has_invalid_cutout_studio_field(image: Image.Image) -> bool:
     """Reject model scenes before alpha extraction when their outer field is not light and neutral.
 
@@ -2535,11 +2621,13 @@ def generate_asset(spec: AssetSpec, index: int, run_id: int) -> tuple[str, str,
     if USE_PRIMARY_IMAGE_MODEL:
         png_content, error = primary_diffusion_png(spec, index, run_id)
         source = f"{PRIMARY_IMAGE_MODEL} + {PRIMARY_LORA_MODEL} on ZeroGPU"
+        if not is_background:
+            source += f" + {FOREGROUND_MODEL_ID} neural alpha"
         if error:
             errors.append(f"primary ZeroGPU diffusion: {error}")
             if REQUIRE_PRIMARY_IMAGE_MODEL:
                 raise RuntimeError(
-                    "Neural image generation is required, but the primary model failed. "
+                    "The neural image pipeline is required, but diffusion generation or foreground extraction failed. "
                     f"No procedural image was returned. Details: {error}"
                 )
 
@@ -2785,7 +2873,7 @@ def model_configuration_summary() -> str:
     sprite_chain = []
     primary_source = f"{PRIMARY_IMAGE_MODEL} + {PRIMARY_LORA_MODEL} on ZeroGPU"
     if USE_PRIMARY_IMAGE_MODEL:
-        sprite_chain.append(f"direct text-to-image {primary_source}")
+        sprite_chain.append(f"direct text-to-image {primary_source} → neural alpha {FOREGROUND_MODEL_ID}")
     if USE_DIFFUSION_FOR_SPRITES:
         sprite_chain.append(FREE_IMAGE_MODEL)
     if USE_HF_IMAGE_PROVIDER:
@@ -2816,6 +2904,7 @@ def model_configuration_summary() -> str:
     neural_models = []
     if USE_PRIMARY_IMAGE_MODEL:
         neural_models.append(f"{PRIMARY_IMAGE_MODEL} + {PRIMARY_LORA_MODEL}")
+        neural_models.append(FOREGROUND_MODEL_ID)
     if USE_DIFFUSION_FOR_SPRITES or USE_DIFFUSION_FOR_BACKGROUNDS:
         neural_models.append(FREE_IMAGE_MODEL)
     if USE_CONTROLNET_FOR_BACKGROUNDS:
@@ -2823,10 +2912,10 @@ def model_configuration_summary() -> str:
     neural_models.extend(remote_sources)
     neural_status = ", ".join(dict.fromkeys(neural_models)) if neural_models else "none active—procedural code only"
     enforcement = "required; generation stops on model failure" if REQUIRE_PRIMARY_IMAGE_MODEL else "development fallback allowed"
-    if PRIMARY_MODEL_ERROR:
-        readiness = f"startup error: {PRIMARY_MODEL_ERROR}"
-    elif PRIMARY_TEXT_PIPE is not None:
-        readiness = "loaded on ZeroGPU"
+    if PRIMARY_MODEL_ERROR or FOREGROUND_MODEL_ERROR:
+        readiness = f"startup error: {PRIMARY_MODEL_ERROR or FOREGROUND_MODEL_ERROR}"
+    elif PRIMARY_TEXT_PIPE is not None and FOREGROUND_MODEL is not None:
+        readiness = "diffusion and neural foreground models loaded on ZeroGPU"
     else:
         readiness = "loads in the deployed ZeroGPU runtime"
     return (
@@ -2837,7 +2926,7 @@ def model_configuration_summary() -> str:
         f"neural enforcement: `{enforcement}` · primary readiness: `{readiness}`. "
         "Every production image is generated directly from its written prompt by the primary text-to-image model. "
         "The user-defined composition and camera contract is placed first inside a conservative CLIP prompt budget. "
-        "Transparent single-subject attempts must first honor the light studio-field contract, are normalized to the requested canvas, then pass configured subject-count, "
+        f"Transparent single-subject attempts use `{FOREGROUND_MODEL_ID}` for neural alpha extraction, are normalized to the requested canvas, then pass configured subject-count, "
         "retained-scene, silhouette, and scale validation; failures are regenerated with a new seed up to "
         f"{PRIMARY_SPRITE_ATTEMPTS} times. No procedural guide is supplied to the model. The procedural renderer is "
         "development-only and is blocked in the deployed Space when the primary model fails."

requirements.txt

@@ -1,8 +1,13 @@
 Pillow
 huggingface_hub
 torch==2.8.0
+torchvision==0.23.0
 diffusers>=0.35.1,<0.38
 transformers>=4.48,<6
 accelerate>=1.2,<2
 peft>=0.14,<1
 safetensors>=0.4
+numpy<3
+timm>=1.0,<2
+kornia>=0.8,<1
+einops>=0.8,<1