Upload pipeline.py with huggingface_hub

pipeline.py

@@ -0,0 +1,373 @@
+"""
+Mongle Character LoRA — Photo-to-Pixel-Art Pipeline
+Standalone script: works after snapshot_download from HuggingFace.
+
+Usage:
+  from huggingface_hub import snapshot_download
+  repo_dir = snapshot_download("Hadimeeee/mongle-character-lora")
+  import sys; sys.path.insert(0, repo_dir)
+  from pipeline import run_pipeline
+  from PIL import Image
+
+  result = run_pipeline(Image.open("photo.jpg"))
+  result["result_nobg"].save("character.png")
+"""
+
+from __future__ import annotations
+
+import os
+import gc
+import json
+import re
+from pathlib import Path
+from typing import Optional
+
+import cv2
+import numpy as np
+import torch
+from PIL import Image
+
+REPO_ID = "Hadimeeee/mongle-character-lora"
+LORA_DIR = Path(__file__).parent   # same folder as this script
+
+# ──────────────────────────────────────────────
+# Image utilities
+# ──────────────────────────────────────────────
+
+def make_square(img: Image.Image, size: int = 1024) -> Image.Image:
+    img = img.convert("RGB")
+    w, h = img.size
+    side = max(w, h)
+    sq = Image.new("RGB", (side, side), (255, 255, 255))
+    sq.paste(img, ((side - w) // 2, (side - h) // 2))
+    return sq.resize((size, size), Image.LANCZOS)
+
+
+def remove_bg(img: Image.Image) -> Image.Image:
+    from rembg import remove as rembg_remove
+    rgba = rembg_remove(img.convert("RGBA"))
+    white = Image.new("RGB", rgba.size, (255, 255, 255))
+    white.paste(rgba, mask=rgba.split()[3])
+    return white
+
+
+def remove_bg_rgba(img: Image.Image) -> Image.Image:
+    from rembg import remove as rembg_remove
+    return rembg_remove(img.convert("RGBA"))
+
+
+# ──────────────────────────────────────────────
+# SAM → flat color → Canny
+# ──────────────────────────────────────────────
+
+def run_sam(img: Image.Image, sam_model: str = "facebook/sam-vit-base"):
+    from transformers import SamModel, SamProcessor
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    processor = SamProcessor.from_pretrained(sam_model)
+    model = SamModel.from_pretrained(sam_model).to(device)
+    model.eval()
+
+    w, h = img.size
+    cx, cy = w // 2, h // 2
+    inputs = processor(img, input_points=[[[cx, cy]]], return_tensors="pt").to(device)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    masks = processor.post_process_masks(
+        outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(),
+        inputs["reshaped_input_sizes"].cpu(),
+    )[0]
+    scores = outputs.iou_scores[0, 0].cpu().numpy()
+    mask = masks[0, int(np.argmax(scores))].numpy().astype(np.uint8) * 255
+    del model, processor; gc.collect(); torch.cuda.empty_cache()
+    return Image.fromarray(mask)
+
+
+def dominant_color(img: Image.Image, mask: Image.Image):
+    arr = np.array(img.convert("RGB"))
+    m   = np.array(mask) > 128
+    px  = arr[m]
+    if len(px) == 0:
+        return (200, 200, 200)
+    from sklearn.cluster import KMeans
+    k = KMeans(n_clusters=3, n_init=5, random_state=0).fit(px)
+    sizes = np.bincount(k.labels_)
+    return tuple(int(c) for c in k.cluster_centers_[np.argmax(sizes)])
+
+
+def build_flat_color(img: Image.Image, mask: Image.Image) -> Image.Image:
+    color = dominant_color(img, mask)
+    flat  = Image.new("RGB", img.size, (255, 255, 255))
+    mask_arr = np.array(mask) > 128
+    flat_arr = np.array(flat)
+    flat_arr[mask_arr] = color
+    return Image.fromarray(flat_arr)
+
+
+def extract_canny(flat: Image.Image, lo: int = 50, hi: int = 150) -> Image.Image:
+    gray  = cv2.cvtColor(np.array(flat), cv2.COLOR_RGB2GRAY)
+    edges = cv2.Canny(gray, lo, hi)
+    return Image.fromarray(np.stack([edges] * 3, axis=-1))
+
+
+# ──────────────────────────────────────────────
+# VLM (Qwen2-VL) — appearance extraction
+# ──────────────────────────────────────────────
+
+_vlm_model = None
+_vlm_proc  = None
+
+
+def load_vlm(model_name: str = "Qwen/Qwen2-VL-7B-Instruct"):
+    global _vlm_model, _vlm_proc
+    from transformers import Qwen2VLForConditionalGeneration, AutoProcessor, BitsAndBytesConfig
+    bnb = BitsAndBytesConfig(load_in_8bit=True)
+    _vlm_model = Qwen2VLForConditionalGeneration.from_pretrained(
+        model_name, quantization_config=bnb, device_map="auto"
+    )
+    _vlm_proc = AutoProcessor.from_pretrained(model_name)
+
+
+def unload_vlm():
+    global _vlm_model, _vlm_proc
+    del _vlm_model, _vlm_proc
+    _vlm_model = _vlm_proc = None
+    gc.collect(); torch.cuda.empty_cache()
+
+
+def run_vlm(img: Image.Image) -> dict:
+    system = (
+        "You are a visual analysis assistant. "
+        "Analyze the stuffed animal in the image and return ONLY a JSON object "
+        "with these fields: animal_type, body_color, secondary_colors (list), "
+        "body_shape, eye_style, accessories (list), distinctive_features (list), "
+        "controlnet_scale (float 0.45-0.85). "
+        "controlnet_scale: 0.45 if no face, 0.5 if pillow-shaped, "
+        "0.75 for normal, 0.85 for limbless/round. "
+        "No explanation, no markdown, only JSON."
+    )
+    messages = [{"role": "user", "content": [
+        {"type": "image", "image": img},
+        {"type": "text",  "text": "Analyze this stuffed animal and return JSON."},
+    ]}]
+    from qwen_vl_utils import process_vision_info
+    text = _vlm_proc.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    image_inputs, _ = process_vision_info(messages)
+    inputs = _vlm_proc(text=[text], images=image_inputs, return_tensors="pt")
+    inputs = {k: v.to(_vlm_model.device) for k, v in inputs.items()}
+    with torch.no_grad():
+        out = _vlm_model.generate(**inputs, max_new_tokens=512, temperature=0.1)
+    raw = _vlm_proc.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
+    m = re.search(r"\{.*\}", raw, re.DOTALL)
+    return json.loads(m.group()) if m else {}
+
+
+def vlm_json_to_prompt(data: dict, extra_en: str = "") -> tuple[str, float]:
+    animal    = data.get("animal_type", "plush toy")
+    body_col  = data.get("body_color", "colorful")
+    sec_cols  = ", ".join(data.get("secondary_colors", []))
+    shape     = data.get("body_shape", "round")
+    eyes      = data.get("eye_style", "round eyes")
+    acc       = ", ".join(data.get("accessories", []))
+    feat      = ", ".join(data.get("distinctive_features", []))
+    cn_scale  = float(data.get("controlnet_scale", 0.75))
+
+    parts = [
+        f"monglestyle, {body_col} {animal} plush",
+        shape, eyes,
+    ]
+    if sec_cols: parts.append(sec_cols)
+    if acc:      parts.append(acc)
+    if feat:     parts.append(feat)
+    if extra_en: parts.append(extra_en)
+    parts += [
+        "single stuffed animal toy mascot character, full body, centered",
+        "front view, cute chibi proportions, 32-bit pixel art sprite",
+        "soft pixel shading, clean silhouette, soft brown outline",
+        "pure white background",
+    ]
+    return ", ".join(p for p in parts if p), cn_scale
+
+
+# ──────────────────────────────────────────────
+# ControlNet generation
+# ──────────────────────────────────────────────
+
+_pipe = None
+
+
+def load_pipeline(lcm: bool = True):
+    global _pipe
+    from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel
+    from diffusers.schedulers import LCMScheduler
+
+    cn = ControlNetModel.from_pretrained(
+        "diffusers/controlnet-canny-sdxl-1.0", torch_dtype=torch.float16
+    )
+    _pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-xl-base-1.0",
+        controlnet=cn, torch_dtype=torch.float16,
+    ).to("cuda")
+
+    if lcm:
+        _pipe.load_lora_weights("latent-consistency/lcm-lora-sdxl", adapter_name="lcm")
+        _pipe.load_lora_weights(str(LORA_DIR), adapter_name="style")
+        _pipe.set_adapters(["lcm", "style"], adapter_weights=[1.0, 0.9])
+        _pipe.scheduler = LCMScheduler.from_config(_pipe.scheduler.config)
+    else:
+        _pipe.load_lora_weights(str(LORA_DIR), adapter_name="style")
+        _pipe.set_adapters(["style"], adapter_weights=[0.9])
+
+
+def unload_pipeline():
+    global _pipe
+    del _pipe; _pipe = None
+    gc.collect(); torch.cuda.empty_cache()
+
+
+def generate_character(
+    canny_img: Image.Image,
+    prompt: str,
+    cn_scale: float = 0.75,
+    steps: int = 8,
+    guidance: float = 1.5,
+    seed: int = 42,
+    lora_scale: float = 0.9,
+) -> Image.Image:
+    neg = "blurry, watermark, text, low quality, deformed, realistic photo, 3d render"
+    gen = torch.Generator("cuda").manual_seed(seed)
+    out = _pipe(
+        prompt=prompt,
+        negative_prompt=neg,
+        image=canny_img,
+        num_inference_steps=steps,
+        guidance_scale=guidance,
+        controlnet_conditioning_scale=cn_scale,
+        cross_attention_kwargs={"scale": lora_scale},
+        generator=gen,
+    )
+    return out.images[0]
+
+
+# ──────────────────────────────────────────────
+# Main API
+# ──────────────────────────────────────────────
+
+def run_pipeline(
+    image_pil: Image.Image,
+    char_desc_en: str = None,
+    lcm: bool = True,
+    lora_scale: float = 0.9,
+    cn_scale_override: float = None,
+    steps: int = 8,
+    seed: int = 42,
+    out_dir: str = None,
+    sam_model: str = "facebook/sam-vit-base",
+    vlm_model: str = "Qwen/Qwen2-VL-7B-Instruct",
+) -> dict:
+    """
+    Full photo-to-pixel-art pipeline.
+
+    Args:
+        image_pil       : Input PIL image (stuffed animal photo)
+        char_desc_en    : Optional English description to supplement VLM output
+        lcm             : Use LCM LoRA for fast 8-step generation
+        lora_scale      : Character LoRA weight (default 0.9)
+        cn_scale_override: Override ControlNet scale (None = VLM recommendation)
+        steps           : Inference steps (8 with LCM, 25-30 without)
+        seed            : Random seed
+        out_dir         : Save intermediate outputs here (optional)
+        sam_model       : SAM model ID
+        vlm_model       : Qwen2-VL model ID
+
+    Returns dict with keys:
+        result, result_nobg, canny, flat_color, appearance, prompt, cn_scale
+    """
+    if out_dir:
+        Path(out_dir).mkdir(parents=True, exist_ok=True)
+
+    # STEP 1: Preprocess
+    print("[1/5] Preprocessing...")
+    sq   = make_square(image_pil)
+    nobg = remove_bg(sq)
+
+    # STEP 2: SAM → flat color → Canny
+    print("[2/5] SAM + Canny edge extraction...")
+    mask  = run_sam(nobg, sam_model)
+    flat  = build_flat_color(nobg, mask)
+    canny = extract_canny(flat)
+
+    # STEP 3: VLM appearance analysis
+    print("[3/5] VLM appearance analysis...")
+    load_vlm(vlm_model)
+    appearance = run_vlm(nobg)
+    unload_vlm()
+
+    prompt, cn_scale = vlm_json_to_prompt(appearance, char_desc_en or "")
+    if cn_scale_override is not None:
+        cn_scale = cn_scale_override
+
+    # STEP 4: Generate character
+    print("[4/5] Generating pixel art character...")
+    guidance = 1.5 if lcm else 7.5
+    load_pipeline(lcm=lcm)
+    result = generate_character(
+        canny, prompt, cn_scale=cn_scale,
+        steps=steps, guidance=guidance, seed=seed, lora_scale=lora_scale,
+    )
+    unload_pipeline()
+
+    # STEP 5: Remove background from result
+    print("[5/5] Final background removal...")
+    result_nobg_rgba = remove_bg_rgba(result)
+
+    # Save outputs
+    if out_dir:
+        d = Path(out_dir)
+        nobg.save(d / "nobg.png")
+        flat.save(d / "flat_color.png")
+        canny.save(d / "canny.png")
+        result.save(d / "result.png")
+        result_nobg_rgba.save(d / "result_nobg.png")
+        (d / "appearance.json").write_text(json.dumps(appearance, ensure_ascii=False, indent=2))
+        (d / "prompt.txt").write_text(prompt)
+        print(f"Saved to: {out_dir}")
+
+    return {
+        "result":      result,
+        "result_nobg": result_nobg_rgba,
+        "canny":       canny,
+        "flat_color":  flat,
+        "appearance":  appearance,
+        "prompt":      prompt,
+        "cn_scale":    cn_scale,
+    }
+
+
+# ──────────────────────────────────────────────
+# CLI
+# ──────────────────────────────────────────────
+
+if __name__ == "__main__":
+    import argparse
+    p = argparse.ArgumentParser(description="Mongle character pipeline")
+    p.add_argument("--image",    required=True, help="Input photo path")
+    p.add_argument("--out-dir",  default="output", help="Output directory")
+    p.add_argument("--desc",     default=None,  help="English character description (optional)")
+    p.add_argument("--no-lcm",   dest="lcm", action="store_false", default=True)
+    p.add_argument("--cn-scale", type=float, default=None)
+    p.add_argument("--steps",    type=int,   default=8)
+    p.add_argument("--seed",     type=int,   default=42)
+    args = p.parse_args()
+
+    result = run_pipeline(
+        image_pil       = Image.open(args.image),
+        char_desc_en    = args.desc,
+        lcm             = args.lcm,
+        cn_scale_override = args.cn_scale,
+        steps           = args.steps,
+        seed            = args.seed,
+        out_dir         = args.out_dir,
+    )
+    print(f"\nPrompt: {result['prompt']}")
+    print(f"cn_scale: {result['cn_scale']}")
+    print(f"Done → {args.out_dir}/result_nobg.png")