Upload demo/real_world_pipeline.py with huggingface_hub

demo/real_world_pipeline.py

@@ -0,0 +1,453 @@
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import re
+import sys
+import time
+import zipfile
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image, ImageOps
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from demo.infer.run_caption_bbox_infer import (  # noqa: E402
+    CAPTION_BBOX_PROMPT_TOP_LEFT,
+    DEFAULT_BBOX_MODEL,
+    draw_boxes,
+    infer_caption_bbox,
+)
+from demo.hf_repo_assets import build_repo_asset_overrides  # noqa: E402
+from demo.infer.vlm_bbox_inference import get_model_and_processor  # noqa: E402
+from infer.common_infer import initialize_pipeline  # noqa: E402
+from infer.infer import build_run_save_dir, get_real_boxes, load_adapter_image  # noqa: E402
+from tools.tools import load_config, seed_everything  # noqa: E402
+
+
+DEFAULT_REAL_CONFIG_PATH = PROJECT_ROOT / "infer" / "infer.yaml"
+DEFAULT_WORK_DIR = PROJECT_ROOT / "demo" / "outputs" / "real_world_demo"
+DEFAULT_RUN_NAME = "step_120000"
+DEFAULT_TARGET_SIZE = 1024
+
+_BBOX_CACHE: dict[str, object] = {"model_path": None, "model": None, "processor": None}
+_REAL_CACHE: dict[str, object] = {"key": None, "pipeline": None, "transp_vae": None}
+
+
+def slugify(text: str) -> str:
+    value = re.sub(r"[^A-Za-z0-9._-]+", "_", text).strip("._-")
+    return value or "sample"
+
+
+def resolve_existing_path(*candidates) -> str | None:
+    for candidate in candidates:
+        if not candidate:
+            continue
+        path = Path(candidate)
+        if path.exists():
+            return str(path)
+    return None
+
+
+DEFAULT_DECOMP_CKPT_ROOT = Path(
+    resolve_existing_path(
+        os.environ.get("SYNLAYERS_DECOMP_CKPT_ROOT"),
+        PROJECT_ROOT / "SynLayers_ckpt" / "step_120000",
+        "/project/llmsvgen/share/data/kmw_layered_checkpoint/SynLayers_ckpt/step_120000",
+    )
+    or "/project/llmsvgen/share/data/kmw_layered_checkpoint/SynLayers_ckpt/step_120000"
+)
+
+
+def prepare_input_image(input_path: str | Path, output_path: Path, size: int) -> Path:
+    image = Image.open(input_path).convert("RGB")
+
+    if image.size != (size, size):
+        resized = ImageOps.contain(image, (size, size), Image.LANCZOS)
+        canvas = Image.new("RGB", (size, size), (255, 255, 255))
+        offset = ((size - resized.width) // 2, (size - resized.height) // 2)
+        canvas.paste(resized, offset)
+        image = canvas
+
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    image.save(output_path)
+    return output_path
+
+
+def load_bbox_bundle(model_path: str):
+    cached_model_path = _BBOX_CACHE["model_path"]
+    if cached_model_path == model_path and _BBOX_CACHE["model"] is not None:
+        return _BBOX_CACHE["model"], _BBOX_CACHE["processor"]
+
+    model, processor = get_model_and_processor(model_path)
+    _BBOX_CACHE.update(
+        {
+            "model_path": model_path,
+            "model": model,
+            "processor": processor,
+        }
+    )
+    return model, processor
+
+
+def load_real_bundle(config: dict):
+    key = (
+        config.get("pretrained_model_name_or_path"),
+        config.get("pretrained_adapter_path"),
+        config.get("transp_vae_path"),
+        config.get("pretrained_lora_dir"),
+        config.get("artplus_lora_dir"),
+        config.get("lora_ckpt"),
+        config.get("layer_ckpt"),
+        config.get("adapter_lora_dir"),
+        config.get("max_layer_num"),
+    )
+
+    if _REAL_CACHE["key"] == key and _REAL_CACHE["pipeline"] is not None:
+        return _REAL_CACHE["pipeline"], _REAL_CACHE["transp_vae"]
+
+    if _REAL_CACHE["pipeline"] is not None:
+        del _REAL_CACHE["pipeline"]
+        del _REAL_CACHE["transp_vae"]
+        _REAL_CACHE["pipeline"] = None
+        _REAL_CACHE["transp_vae"] = None
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    pipeline, transp_vae = initialize_pipeline(config)
+    _REAL_CACHE.update({"key": key, "pipeline": pipeline, "transp_vae": transp_vae})
+    return pipeline, transp_vae
+
+
+def build_runtime_config(
+    *,
+    config_path: str | Path,
+    image_dir: Path,
+    bbox_jsonl: Path,
+    results_root: Path,
+    run_name: str,
+    seed: int | None = None,
+) -> dict:
+    config = load_config(str(config_path))
+    repo_overrides = build_repo_asset_overrides(os.environ.get("SYNLAYERS_MODEL_REPO"))
+    decomp_ckpt_root = Path(
+        os.environ.get("SYNLAYERS_DECOMP_CKPT_ROOT")
+        or repo_overrides.get("decomp_ckpt_root")
+        or DEFAULT_DECOMP_CKPT_ROOT
+    )
+    config["data_dir"] = str(image_dir.parent)
+    config["image_dir"] = str(image_dir)
+    config["test_jsonl"] = str(bbox_jsonl)
+    config["save_dir"] = str(results_root)
+    config["run_name"] = run_name
+    config["lora_ckpt"] = str(decomp_ckpt_root / "transformer")
+    config["layer_ckpt"] = str(decomp_ckpt_root)
+    config["adapter_lora_dir"] = str(decomp_ckpt_root / "adapter")
+
+    env_overrides = {
+        "pretrained_model_name_or_path": (
+            os.environ.get("SYNLAYERS_BASE_MODEL")
+            or repo_overrides.get("pretrained_model_name_or_path")
+            or resolve_existing_path(PROJECT_ROOT / "SynLayers_checkpoints" / "FLUX.1-dev")
+            or "black-forest-labs/FLUX.1-dev"
+        ),
+        "pretrained_adapter_path": (
+            os.environ.get("SYNLAYERS_ADAPTER_MODEL")
+            or repo_overrides.get("pretrained_adapter_path")
+            or resolve_existing_path(
+                PROJECT_ROOT / "SynLayers_checkpoints" / "FLUX.1-dev-Controlnet-Inpainting-Alpha"
+            )
+        ),
+        "transp_vae_path": (
+            os.environ.get("SYNLAYERS_TRANSP_VAE")
+            or repo_overrides.get("transp_vae_path")
+            or resolve_existing_path(PROJECT_ROOT / "ckpt" / "trans_vae" / "0008000.pt")
+        ),
+        "pretrained_lora_dir": (
+            os.environ.get("SYNLAYERS_PRETRAINED_LORA")
+            or repo_overrides.get("pretrained_lora_dir")
+            or resolve_existing_path(PROJECT_ROOT / "ckpt" / "pre_trained_LoRA")
+        ),
+        "artplus_lora_dir": (
+            os.environ.get("SYNLAYERS_ARTPLUS_LORA")
+            or repo_overrides.get("artplus_lora_dir")
+            or resolve_existing_path(PROJECT_ROOT / "ckpt" / "prism_ft_LoRA")
+        ),
+    }
+    for key, value in env_overrides.items():
+        if value:
+            config[key] = value
+
+    if seed is not None:
+        config["seed"] = seed
+
+    return config
+
+
+def write_bbox_jsonl(record: dict, output_path: Path) -> Path:
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with output_path.open("w", encoding="utf-8") as handle:
+        handle.write(json.dumps(record, ensure_ascii=False) + "\n")
+    return output_path
+
+
+def format_source_image_path(image_path: str, image_dir: Path) -> str:
+    path = Path(image_path)
+    try:
+        return path.relative_to(image_dir).as_posix()
+    except ValueError:
+        return path.name
+
+
+def save_real_case(
+    *,
+    sample: dict,
+    config: dict,
+    pipeline,
+    transp_vae,
+) -> dict:
+    if config.get("seed") is not None:
+        seed_everything(config["seed"])
+
+    source_size = config.get("source_size", DEFAULT_TARGET_SIZE)
+    target_size = config.get("target_size", DEFAULT_TARGET_SIZE)
+    max_layer_num = config.get("max_layer_num", 52)
+    sample_name = sample["sample_or_stem"]
+
+    layer_boxes = get_real_boxes(sample, source_size, target_size)
+    adapter_img, resolved_image_path = load_adapter_image(sample, target_size, config)
+
+    whole_box = (0, 0, target_size, target_size)
+    bg_box = (0, 0, target_size, target_size)
+    all_boxes = [whole_box, bg_box] + layer_boxes
+    if len(all_boxes) > max_layer_num:
+        raise ValueError(
+            f"num_layers={len(all_boxes)} exceeds max_layer_num={max_layer_num} for {sample_name}"
+        )
+
+    generator = torch.Generator(device=torch.device("cuda")).manual_seed(config.get("seed", 42))
+    caption = sample.get("whole_caption", "")
+
+    x_hat, image, _ = pipeline(
+        prompt=caption,
+        adapter_image=adapter_img,
+        adapter_conditioning_scale=config.get("adapter_scale", 0.9),
+        validation_box=all_boxes,
+        generator=generator,
+        height=target_size,
+        width=target_size,
+        guidance_scale=config.get("cfg", 4.0),
+        num_layers=len(all_boxes),
+        sdxl_vae=transp_vae,
+    )
+
+    x_hat = (x_hat + 1) / 2
+    x_hat = x_hat.squeeze(0).permute(1, 0, 2, 3).to(torch.float32)
+
+    save_dir, resolved_run_name = build_run_save_dir(config)
+    save_dir_path = Path(save_dir)
+    case_dir = save_dir_path / sample_name
+    merged_dir = save_dir_path / "merged"
+    merged_rgba_dir = save_dir_path / "merged_rgba"
+    case_dir.mkdir(parents=True, exist_ok=True)
+    merged_dir.mkdir(parents=True, exist_ok=True)
+    merged_rgba_dir.mkdir(parents=True, exist_ok=True)
+
+    whole_rgba_path = case_dir / "whole_image_rgba.png"
+    background_rgba_path = case_dir / "background_rgba.png"
+    origin_path = case_dir / "origin.png"
+    merged_case_path = case_dir / "merged.png"
+    merged_global_path = merged_dir / f"{sample_name}.png"
+    merged_rgba_path = merged_rgba_dir / f"{sample_name}.png"
+
+    whole_image_layer = (x_hat[0].permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
+    Image.fromarray(whole_image_layer, "RGBA").save(whole_rgba_path)
+
+    background_layer = (x_hat[1].permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
+    Image.fromarray(background_layer, "RGBA").save(background_rgba_path)
+
+    adapter_img.save(origin_path)
+
+    merged_image = image[1]
+    layer_paths: list[str] = []
+    for layer_idx in range(2, x_hat.shape[0]):
+        rgba_layer = (x_hat[layer_idx].permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
+        rgba_image = Image.fromarray(rgba_layer, "RGBA")
+        layer_path = case_dir / f"layer_{layer_idx - 2}_rgba.png"
+        rgba_image.save(layer_path)
+        layer_paths.append(str(layer_path))
+        merged_image = Image.alpha_composite(merged_image.convert("RGBA"), rgba_image)
+
+    merged_image.convert("RGB").save(merged_global_path)
+    merged_image.convert("RGB").save(merged_case_path)
+    merged_image.save(merged_rgba_path)
+
+    case_meta = {
+        "sample_name": sample_name,
+        "source_image_path": format_source_image_path(
+            resolved_image_path,
+            Path(config["image_dir"]),
+        ),
+        "target_size": target_size,
+        "source_size": source_size,
+        "raw_num_layers": sample.get("num_layers"),
+        "num_layers": len(all_boxes),
+        "raw_boxes": sample.get("bboxes", []),
+        "boxes": all_boxes,
+        "caption": caption,
+        "run_name": resolved_run_name,
+    }
+    meta_path = case_dir / "inference_meta.json"
+    with meta_path.open("w", encoding="utf-8") as handle:
+        json.dump(case_meta, handle, indent=2)
+
+    return {
+        "run_name": resolved_run_name,
+        "save_dir": str(save_dir_path),
+        "case_dir": str(case_dir),
+        "merged_image": str(merged_case_path),
+        "merged_global_image": str(merged_global_path),
+        "merged_rgba_image": str(merged_rgba_path),
+        "whole_image_rgba": str(whole_rgba_path),
+        "background_rgba": str(background_rgba_path),
+        "origin_image": str(origin_path),
+        "layer_images": layer_paths,
+        "metadata_path": str(meta_path),
+        "metadata": case_meta,
+    }
+
+
+def create_archive(run_dir: Path) -> Path:
+    archive_path = run_dir / "synlayers_result_bundle.zip"
+    with zipfile.ZipFile(archive_path, "w", compression=zipfile.ZIP_DEFLATED) as zf:
+        for path in run_dir.rglob("*"):
+            if path == archive_path or path.is_dir():
+                continue
+            zf.write(path, arcname=path.relative_to(run_dir))
+    return archive_path
+
+
+def run_real_world_pipeline(
+    image_path: str | Path,
+    *,
+    sample_name: str | None = None,
+    work_dir: str | Path | None = None,
+    bbox_model: str | None = None,
+    config_path: str | Path | None = None,
+    max_new_tokens: int = 1024,
+    seed: int | None = None,
+    run_name: str = DEFAULT_RUN_NAME,
+) -> dict:
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA GPU is required for the unified SynLayers real-world pipeline.")
+
+    image_path = Path(image_path)
+    if not image_path.exists():
+        raise FileNotFoundError(f"Input image not found: {image_path}")
+
+    bbox_model = (
+        bbox_model
+        or os.environ.get("SYNLAYERS_BBOX_MODEL")
+        or os.environ.get("SYNLAYERS_MODEL_REPO")
+        or DEFAULT_BBOX_MODEL
+    )
+    config_path = Path(config_path or os.environ.get("SYNLAYERS_REAL_CONFIG", str(DEFAULT_REAL_CONFIG_PATH)))
+    work_dir = Path(work_dir or os.environ.get("SYNLAYERS_DEMO_WORK_DIR", str(DEFAULT_WORK_DIR)))
+
+    normalized_sample_name = slugify(sample_name or image_path.stem)
+    timestamp = f"{time.strftime('%Y%m%d_%H%M%S')}_{int((time.time() % 1) * 1000):03d}"
+    run_dir = work_dir / f"{timestamp}_{normalized_sample_name}"
+    image_dir = run_dir / "layers_real_test_1024"
+    prepared_image_path = prepare_input_image(
+        image_path,
+        image_dir / f"{normalized_sample_name}.png",
+        DEFAULT_TARGET_SIZE,
+    )
+
+    bbox_model_bundle, bbox_processor = load_bbox_bundle(bbox_model)
+    whole_caption, bboxes = infer_caption_bbox(
+        prepared_image_path,
+        bbox_model_bundle,
+        bbox_processor,
+        prompt=CAPTION_BBOX_PROMPT_TOP_LEFT,
+        max_new_tokens=max_new_tokens,
+    )
+
+    record = {
+        "sample_or_stem": normalized_sample_name,
+        "image": prepared_image_path.name,
+        "whole_caption": whole_caption,
+        "bboxes": bboxes,
+        "num_layers": len(bboxes),
+        "coord": "top_left",
+    }
+
+    bbox_jsonl = write_bbox_jsonl(record, run_dir / "caption_bbox_infer.jsonl")
+    bbox_vis_path = run_dir / "bbox_vis" / f"{normalized_sample_name}_vis.png"
+    draw_boxes(prepared_image_path, bboxes, bbox_vis_path)
+
+    config = build_runtime_config(
+        config_path=config_path,
+        image_dir=image_dir,
+        bbox_jsonl=bbox_jsonl,
+        results_root=run_dir / "results",
+        run_name=run_name,
+        seed=seed,
+    )
+    pipeline, transp_vae = load_real_bundle(config)
+    decomposition_result = save_real_case(
+        sample=record,
+        config=config,
+        pipeline=pipeline,
+        transp_vae=transp_vae,
+    )
+
+    archive_path = create_archive(run_dir)
+    decomposition_result.update(
+        {
+            "input_image": str(prepared_image_path),
+            "bbox_visualization": str(bbox_vis_path),
+            "bbox_jsonl": str(bbox_jsonl),
+            "bbox_record": record,
+            "archive_path": str(archive_path),
+            "config_path": str(config_path),
+            "bbox_model": bbox_model,
+        }
+    )
+    return decomposition_result
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Run the unified real-world SynLayers pipeline on one image."
+    )
+    parser.add_argument("--image", type=str, required=True, help="Input image path")
+    parser.add_argument("--sample-name", type=str, default=None)
+    parser.add_argument("--work-dir", type=str, default=str(DEFAULT_WORK_DIR))
+    parser.add_argument("--bbox-model", type=str, default=DEFAULT_BBOX_MODEL)
+    parser.add_argument("--config", type=str, default=str(DEFAULT_REAL_CONFIG_PATH))
+    parser.add_argument("--max-new-tokens", type=int, default=1024)
+    parser.add_argument("--seed", type=int, default=None)
+    parser.add_argument("--run-name", type=str, default=DEFAULT_RUN_NAME)
+    args = parser.parse_args()
+
+    result = run_real_world_pipeline(
+        args.image,
+        sample_name=args.sample_name,
+        work_dir=args.work_dir,
+        bbox_model=args.bbox_model,
+        config_path=args.config,
+        max_new_tokens=args.max_new_tokens,
+        seed=args.seed,
+        run_name=args.run_name,
+    )
+    print(json.dumps(result, indent=2, ensure_ascii=False))
+
+
+if __name__ == "__main__":
+    main()