| """NitroFusion2 text-to-image with diffusers (1-4 steps). |
| |
| python inference.py --steps 2 --out example.png |
| """ |
| import argparse |
|
|
| import torch |
| from diffusers import FlowMatchEulerDiscreteScheduler, ZImagePipeline, ZImageTransformer2DModel |
|
|
| BASE = "Tongyi-MAI/Z-Image-Turbo" |
| REPO = "ChenDY/NitroFusion2" |
| DEDISTILL_LORA = "nf2-dedistill-lora.safetensors" |
| SHIFT = {1: 3.0, 2: 6.0, 3: 8.0, 4: 9.0} |
| LORA_SCALES = {1: (0.0,), 2: (0.5, 0.0), 3: (0.5, 0.6, 0.0), 4: (0.5, 0.6, 0.7, 0.0)} |
|
|
|
|
| def load(repo: str = REPO, dedistill_lora: bool = True) -> ZImagePipeline: |
| transformer = ZImageTransformer2DModel.from_pretrained(repo, subfolder="transformer", torch_dtype=torch.bfloat16) |
| pipe = ZImagePipeline.from_pretrained(BASE, transformer=transformer, torch_dtype=torch.bfloat16) |
| if dedistill_lora: |
| pipe.load_lora_weights(repo, weight_name=DEDISTILL_LORA, adapter_name="dedistill") |
| return pipe |
|
|
|
|
| @torch.no_grad() |
| def generate(pipe: ZImagePipeline, prompt: str, steps: int = 2, height: int = 1024, width: int = 1024, |
| seed: int = 0, dedistill_lora: bool = True): |
| assert steps in SHIFT, "NitroFusion2 is distilled for 1-4 steps" |
| pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(pipe.scheduler.config, shift=SHIFT[steps]) |
| sigmas = [1.0 - i / steps for i in range(steps)] |
| scales = LORA_SCALES[steps] if dedistill_lora else (0.0,) * steps |
| callback = None |
| if dedistill_lora: |
| pipe.set_adapters("dedistill", scales[0]) |
|
|
| def callback(p, i, t, kwargs): |
| if i + 1 < steps: |
| p.set_adapters("dedistill", scales[i + 1]) |
| return kwargs |
|
|
| image = pipe(prompt=prompt, height=height, width=width, num_inference_steps=steps, sigmas=sigmas, |
| guidance_scale=0.0, generator=torch.Generator("cuda").manual_seed(seed), |
| callback_on_step_end=callback).images[0] |
| return image |
|
|
|
|
| if __name__ == "__main__": |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--prompt", default='Latina female with thick wavy hair, harbor boats and pastel houses behind. Breezy seaside light, warm tones, cinematic close-up.') |
| ap.add_argument("--steps", type=int, default=2) |
| ap.add_argument("--height", type=int, default=1024) |
| ap.add_argument("--width", type=int, default=1024) |
| ap.add_argument("--seed", type=int, default=0) |
| ap.add_argument("--no-dedistill-lora", action="store_true") |
| ap.add_argument("--repo", default=REPO) |
| ap.add_argument("--out", default="nitrofusion2.png") |
| args = ap.parse_args() |
| pipe = load(args.repo, dedistill_lora=not args.no_dedistill_lora) |
| pipe.enable_model_cpu_offload() |
| generate(pipe, args.prompt, args.steps, args.height, args.width, args.seed, dedistill_lora=not args.no_dedistill_lora).save(args.out) |
| print("saved", args.out) |
|
|
