Add vnasset edit command

Qwen Image Edit pipeline with FP8 Safetensors loading.
Uses init_empty_weights for memory-efficient 40GB model loading.
bf16 dtype to avoid ROCm crashes; falls back to math SDPA.

vnassets/cli.py

@@ -1,6 +1,7 @@
 """VNAsset — Visual Novel Asset Pipeline CLI."""
 import click
 
+from .edit import edit
 from .generate import generate
 
 
@@ -46,3 +47,29 @@ def generate_cmd(checkpoint, prompt, negative_prompt, width, height, steps, cfg,
         seed=seed,
         output_path=output,
     )
+
+
+@main.command()
+@click.option("--model", required=True, help="Path to Qwen Image Edit diffusion model (.safetensors)")
+@click.option("--input", "input_path", required=True, help="Input image to edit")
+@click.option("--prompt", required=True, help="Edit instruction")
+@click.option("--steps", default=20, type=int, help="Inference steps (4 for turbo)")
+@click.option("--cfg", default=4.0, type=float, help="CFG scale (1.0 for turbo)")
+@click.option(
+    "--seed", default="random", callback=_parse_seed,
+    help="RNG seed (integer or 'random')",
+)
+@click.option("--lora", "lora_path", default=None, help="Path to LoRA weights (.safetensors)")
+@click.option("--output", default="output.png", help="Output image path")
+def edit_cmd(model, input_path, prompt, steps, cfg, seed, lora_path, output):
+    """Edit an image using Qwen Image Edit."""
+    edit(
+        model_path=model,
+        input_path=input_path,
+        prompt=prompt,
+        steps=steps,
+        cfg=cfg,
+        seed=seed,
+        lora_path=lora_path,
+        output_path=output,
+    )

vnassets/edit.py

@@ -0,0 +1,131 @@
+"""Qwen Image Edit — image-to-image editing."""
+import gc
+import json
+import random
+import time
+from pathlib import Path
+
+import safetensors.torch
+import torch
+from accelerate import init_empty_weights
+from PIL import Image
+from diffusers import QwenImageEditPlusPipeline, FlowMatchEulerDiscreteScheduler
+from diffusers.models.autoencoders import AutoencoderKLQwenImage
+from diffusers.models.transformers import QwenImageTransformer2DModel
+from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessor
+
+
+TEXT_ENCODER_ID = "Qwen/Qwen2.5-VL-7B-Instruct"
+VAE_ID = "Qwen/Qwen-Image"
+
+
+def _load_transformer(path: str, dtype: torch.dtype) -> QwenImageTransformer2DModel:
+    """Load Qwen Image Edit transformer from a single FP8 safetensors file.
+
+    Uses init_empty_weights and incremental conversion to keep peak memory
+    manageable. The model is 20B parameters (20 GB FP8, 40 GB BF16)."""
+    config = QwenImageTransformer2DModel.load_config(
+        "Qwen/Qwen-Image-Edit", subfolder="transformer"
+    )
+
+    state_dict = safetensors.torch.load_file(path)
+    prefix = "model.diffusion_model."
+
+    # Convert FP8 -> target dtype, freeing FP8 tensors as we go
+    cleaned = {}
+    for k in list(state_dict.keys()):
+        if k.startswith(prefix):
+            v = state_dict.pop(k)
+            cleaned[k[len(prefix):]] = v.to(dtype)
+            del v
+    del state_dict
+    gc.collect()
+
+    # Create model on meta device to avoid allocating full model in addition to cleaned dict
+    with init_empty_weights():
+        model = QwenImageTransformer2DModel.from_config(config, torch_dtype=dtype)
+
+    model.load_state_dict(cleaned, strict=True, assign=True)
+    del cleaned
+    gc.collect()
+    return model
+
+
+def edit(
+    model_path: str,
+    input_path: str,
+    prompt: str,
+    steps: int = 20,
+    cfg: float = 4.0,
+    seed: int | None = None,
+    lora_path: str | None = None,
+    output_path: str = "output.png",
+) -> None:
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = torch.bfloat16
+
+    if seed is None:
+        seed = random.randint(0, 2**32 - 1)
+
+    output = Path(output_path)
+    output.parent.mkdir(parents=True, exist_ok=True)
+
+    t0 = time.perf_counter()
+
+    transformer = _load_transformer(model_path, dtype)
+    vae = AutoencoderKLQwenImage.from_pretrained(VAE_ID, subfolder="vae", torch_dtype=dtype)
+    text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+        TEXT_ENCODER_ID, torch_dtype=dtype
+    )
+    tokenizer = Qwen2Tokenizer.from_pretrained(TEXT_ENCODER_ID)
+    processor = Qwen2VLProcessor.from_pretrained(TEXT_ENCODER_ID)
+    scheduler = FlowMatchEulerDiscreteScheduler()
+
+    pipe = QwenImageEditPlusPipeline(
+        scheduler=scheduler,
+        vae=vae,
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        processor=processor,
+        transformer=transformer,
+    )
+    pipe.to(device)
+
+    if lora_path:
+        pipe.load_lora_weights(lora_path)
+        pipe.fuse_lora()
+
+    t_load = time.perf_counter() - t0
+
+    input_image = Image.open(input_path).convert("RGB")
+    generator = torch.Generator(device=device).manual_seed(seed)
+
+    t1 = time.perf_counter()
+    image = pipe(
+        image=input_image,
+        prompt=prompt,
+        true_cfg_scale=cfg,
+        num_inference_steps=steps,
+        generator=generator,
+    ).images[0]
+    t_infer = time.perf_counter() - t1
+
+    image.save(output)
+    print(f"Saved {output}")
+
+    meta_path = output.with_suffix(".json")
+    meta = {
+        "model": str(Path(model_path).resolve()),
+        "vae": VAE_ID,
+        "text_encoder": TEXT_ENCODER_ID,
+        "input_image": str(Path(input_path).resolve()),
+        "prompt": prompt,
+        "steps": steps,
+        "cfg": cfg,
+        "seed": seed,
+        "lora": str(Path(lora_path).resolve()) if lora_path else None,
+        "load_time_s": round(t_load, 2),
+        "inference_time_s": round(t_infer, 2),
+    }
+    meta_path.write_text(json.dumps(meta, indent=2))
+    print(f"Saved {meta_path}")