File size: 6,083 Bytes
ce47bc4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | #!/usr/bin/env python3
"""Batch image editing with LongCat-Image-Edit and a PEFT LoRA adapter."""
from __future__ import annotations
import argparse
from pathlib import Path
import torch
from diffusers import LongCatImageEditPipeline
from peft import PeftModel
from PIL import Image
IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
DEFAULT_BASE = "meituan-longcat/LongCat-Image-Edit"
DEFAULT_PROMPT = (
"Edit only the human hand region. Replace the human hand with a realistic Inspire robotic "
"hand with correct robotic finger structure and joints. Preserve the original wrist pose, "
"palm orientation, finger articulation, grasp geometry, and contact points with the object. "
"The robot hand must be kinematically feasible and physically plausible, without penetrating "
"the object. Keep the object pose, shape, texture, background, lighting, camera viewpoint, "
"and all non-hand regions unchanged."
)
def list_images(input_dir: Path, recursive: bool) -> list[Path]:
iterator = input_dir.rglob("*") if recursive else input_dir.iterdir()
return sorted(
path
for path in iterator
if path.is_file() and path.suffix.lower() in IMAGE_EXTS
)
def apply_lora_scale(model: torch.nn.Module, scale: float) -> int:
"""Multiply every active PEFT LoRA layer scale by ``scale``."""
if scale == 1.0:
return 0
changed = 0
for module in model.modules():
scaling = getattr(module, "scaling", None)
if isinstance(scaling, dict):
for adapter_name in list(scaling):
scaling[adapter_name] *= scale
changed += 1
return changed
def main() -> None:
parser = argparse.ArgumentParser(
description="Batch inference for LongCat-Image-Edit with the HandEdit LoRA."
)
parser.add_argument(
"--base",
default=DEFAULT_BASE,
help=f"Base model directory or model ID (default: {DEFAULT_BASE}).",
)
parser.add_argument(
"--lora",
default="./checkpoints/longcat",
help="PEFT LoRA directory (default: ./checkpoints/longcat).",
)
parser.add_argument("--input_dir", required=True)
parser.add_argument("--output_dir", required=True)
parser.add_argument("--prompt", default=DEFAULT_PROMPT)
parser.add_argument("--negative_prompt", default="")
parser.add_argument("--guidance_scale", type=float, default=4.5)
parser.add_argument("--steps", type=int, default=50)
parser.add_argument("--seed", type=int, default=43)
parser.add_argument(
"--seed_mode", choices=["fixed", "increment"], default="fixed"
)
parser.add_argument("--lora_scale", type=float, default=1.0)
parser.add_argument(
"--offload", choices=["model", "sequential", "none"], default="model"
)
parser.add_argument("--local_files_only", action="store_true")
parser.add_argument("--suffix", default="")
parser.add_argument("--recursive", action="store_true")
parser.add_argument("--skip_existing", action="store_true")
args = parser.parse_args()
if not torch.cuda.is_available():
raise RuntimeError("CUDA is required for practical LongCat inference.")
if args.steps < 2:
raise ValueError("--steps must be at least 2")
if args.guidance_scale < 0:
raise ValueError("--guidance_scale must be non-negative")
input_root = Path(args.input_dir).expanduser().resolve()
output_root = Path(args.output_dir).expanduser().resolve()
lora_root = Path(args.lora).expanduser().resolve()
if not input_root.is_dir():
raise FileNotFoundError(f"Input directory not found: {input_root}")
if not lora_root.is_dir():
raise FileNotFoundError(f"LoRA directory not found: {lora_root}")
paths = list_images(input_root, args.recursive)
if not paths:
raise RuntimeError(f"No images found under: {input_root}")
output_root.mkdir(parents=True, exist_ok=True)
print("[1/3] Loading LongCat-Image-Edit...")
pipe = LongCatImageEditPipeline.from_pretrained(
args.base,
torch_dtype=torch.bfloat16,
local_files_only=args.local_files_only,
)
print("[2/3] Loading HandEdit LoRA...")
pipe.transformer = PeftModel.from_pretrained(
pipe.transformer,
str(lora_root),
is_trainable=False,
)
changed = apply_lora_scale(pipe.transformer, args.lora_scale)
if args.lora_scale != 1.0:
print(
f"[INFO] LoRA scale={args.lora_scale}; "
f"adjusted {changed} active LoRA layers."
)
if args.offload == "model":
pipe.enable_model_cpu_offload()
elif args.offload == "sequential":
pipe.enable_sequential_cpu_offload()
else:
pipe.to("cuda")
print(f"[3/3] Processing {len(paths)} images...")
current_seed = args.seed
for index, input_path in enumerate(paths, start=1):
relative = input_path.relative_to(input_root)
output_path = output_root / relative.with_name(
relative.stem + args.suffix + ".png"
)
output_path.parent.mkdir(parents=True, exist_ok=True)
if args.skip_existing and output_path.exists():
print(f"[{index}/{len(paths)}] SKIP {output_path}")
else:
with Image.open(input_path) as source:
image = source.convert("RGB")
generator = torch.Generator("cpu").manual_seed(current_seed)
result = pipe(
image,
args.prompt,
negative_prompt=args.negative_prompt,
guidance_scale=args.guidance_scale,
num_inference_steps=args.steps,
num_images_per_prompt=1,
generator=generator,
).images[0]
result.save(output_path)
print(f"[{index}/{len(paths)}] OK {input_path.name} -> {output_path}")
if args.seed_mode == "increment":
current_seed += 1
print(f"[DONE] Results saved under: {output_root}")
if __name__ == "__main__":
main()
|