File size: 2,661 Bytes
76d175c | 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 63 64 65 66 67 68 | """Generate an image with the DiffusionART SD3.5-M multi-reward PEFT adapter."""
import argparse
from pathlib import Path
BASE_MODEL = "stabilityai/stable-diffusion-3.5-medium"
BASE_REVISION = "b940f670f0eda2d07fbb75229e779da1ad11eb80"
def load_pipeline(adapter, *, base_revision=BASE_REVISION, device="cuda", offload=False):
import torch
from diffusers import StableDiffusion3Pipeline
from peft import PeftModel
if not device.startswith("cuda") or not torch.cuda.is_available():
raise RuntimeError("This inference example requires an available CUDA GPU")
pipe = StableDiffusion3Pipeline.from_pretrained(
BASE_MODEL, revision=base_revision, torch_dtype=torch.float16,
)
pipe.transformer = PeftModel.from_pretrained(
pipe.transformer.float(), adapter, is_trainable=False,
adapter_name="diffusionart",
)
pipe.transformer.eval()
pipe.vae.to(dtype=torch.float32)
if offload:
pipe.enable_model_cpu_offload(device=device)
else:
pipe.to(device)
return pipe
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--adapter", default=str(Path(__file__).resolve().parent))
parser.add_argument("--prompt", required=True)
parser.add_argument("--output", default="image.png")
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--steps", type=int, default=40)
parser.add_argument("--height", type=int, default=512)
parser.add_argument("--width", type=int, default=512)
parser.add_argument("--guidance-scale", type=float, default=1.0)
parser.add_argument("--device", default="cuda")
parser.add_argument("--cpu-offload", action="store_true")
args = parser.parse_args()
if args.steps <= 0 or any(n <= 0 or n % 16 for n in (args.height, args.width)):
parser.error("steps must be positive; dimensions must be positive multiples of 16")
output = Path(args.output)
if output.exists():
parser.error("output already exists; choose another path")
import torch
pipe = load_pipeline(args.adapter, device=args.device, offload=args.cpu_offload)
generator = torch.Generator(device="cpu").manual_seed(args.seed)
with torch.inference_mode(), torch.autocast("cuda", dtype=torch.float16):
image = pipe(
args.prompt, num_inference_steps=args.steps, guidance_scale=args.guidance_scale,
height=args.height, width=args.width, generator=generator,
).images[0]
output.parent.mkdir(parents=True, exist_ok=True)
image.save(output)
print(output.resolve())
if __name__ == "__main__":
main()
|