diff --git a/vnassets/cli.py b/vnassets/cli.py index c51e11b..bc0930b 100644 --- a/vnassets/cli.py +++ b/vnassets/cli.py @@ -53,15 +53,14 @@ def generate_cmd(checkpoint, prompt, negative_prompt, width, height, steps, cfg, @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("--steps", default=20, type=int, help="Inference steps") +@click.option("--cfg", default=4.0, type=float, help="CFG scale") @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): +def edit_cmd(model, input_path, prompt, steps, cfg, seed, output): """Edit an image using Qwen Image Edit.""" edit( model_path=model, @@ -70,6 +69,5 @@ def edit_cmd(model, input_path, prompt, steps, cfg, seed, lora_path, output): steps=steps, cfg=cfg, seed=seed, - lora_path=lora_path, output_path=output, ) diff --git a/vnassets/edit.py b/vnassets/edit.py index 2303c81..b11f290 100644 --- a/vnassets/edit.py +++ b/vnassets/edit.py @@ -58,7 +58,6 @@ def edit( 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" @@ -91,13 +90,20 @@ def edit( ) pipe.to(device) - if lora_path: - pipe.load_lora_weights(lora_path) - pipe.fuse_lora() - t_load = time.perf_counter() - t0 + # Resize large inputs to avoid GPU crashes. The Qwen transformer does + # O(N^2) attention on image patches. 512px is safe on this GPU. input_image = Image.open(input_path).convert("RGB") + w, h = input_image.size + max_dim = 512 + if w > max_dim or h > max_dim: + scale = max_dim / max(w, h) + new_w, new_h = int(w * scale), int(h * scale) + new_w = (new_w // 16) * 16 + new_h = (new_h // 16) * 16 + input_image = input_image.resize((new_w, new_h), Image.LANCZOS) + print(f"Resized input {w}x{h} -> {new_w}x{new_h}") generator = torch.Generator(device=device).manual_seed(seed) t1 = time.perf_counter() @@ -113,6 +119,10 @@ def edit( image.save(output) print(f"Saved {output}") + del pipe, transformer, vae, text_encoder + if device == "cuda": + torch.cuda.empty_cache() + meta_path = output.with_suffix(".json") meta = { "model": str(Path(model_path).resolve()), @@ -123,7 +133,6 @@ def edit( "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), } diff --git a/vnassets/generate.py b/vnassets/generate.py index 8055ee9..307d088 100644 --- a/vnassets/generate.py +++ b/vnassets/generate.py @@ -57,6 +57,11 @@ def generate( image.save(output) print(f"Saved {output}") + # Free GPU memory before returning + del pipe + if device == "cuda": + torch.cuda.empty_cache() + meta_path = output.with_suffix(".json") meta = { "checkpoint": str(Path(checkpoint_path).resolve()),