From ed124b1cbba464f440471536f033e02fec9665ec Mon Sep 17 00:00:00 2001 From: Michele Rossi Date: Mon, 6 Jul 2026 18:06:47 +0200 Subject: [PATCH] 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 | 27 ++++++++++ vnassets/edit.py | 131 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 158 insertions(+) create mode 100644 vnassets/edit.py diff --git a/vnassets/cli.py b/vnassets/cli.py index 5ef1bfe..c51e11b 100644 --- a/vnassets/cli.py +++ b/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, + ) diff --git a/vnassets/edit.py b/vnassets/edit.py new file mode 100644 index 0000000..2303c81 --- /dev/null +++ b/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}")