File size: 2,819 Bytes
5d0a7f5 | 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 | """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)
|