06fba9c234
ROCm-safe bfloat16 inference with custom matmul attention. Automatic output directories, random seed, timing metadata.
75 lines
1.9 KiB
Python
75 lines
1.9 KiB
Python
"""SDXL text-to-image generation."""
|
|
import json
|
|
import random
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from diffusers import StableDiffusionXLPipeline
|
|
|
|
from .attention import patch_unet_attention
|
|
|
|
|
|
def generate(
|
|
checkpoint_path: str,
|
|
prompt: str,
|
|
negative_prompt: str = "",
|
|
width: int = 1024,
|
|
height: int = 1024,
|
|
steps: int = 20,
|
|
cfg: float = 4.5,
|
|
seed: int | None = None,
|
|
output_path: str = "output.png",
|
|
) -> None:
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
# bfloat16 avoids ROCm kernel crashes on RDNA 3.5; float16 segfaults
|
|
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()
|
|
pipe = StableDiffusionXLPipeline.from_single_file(
|
|
checkpoint_path,
|
|
torch_dtype=dtype,
|
|
)
|
|
pipe.to(device)
|
|
patch_unet_attention(pipe.unet)
|
|
t_load = time.perf_counter() - t0
|
|
|
|
generator = torch.Generator(device=device).manual_seed(seed)
|
|
|
|
t1 = time.perf_counter()
|
|
image = pipe(
|
|
prompt=prompt,
|
|
negative_prompt=negative_prompt,
|
|
width=width,
|
|
height=height,
|
|
num_inference_steps=steps,
|
|
guidance_scale=cfg,
|
|
generator=generator,
|
|
).images[0]
|
|
t_infer = time.perf_counter() - t1
|
|
|
|
image.save(output)
|
|
print(f"Saved {output}")
|
|
|
|
meta_path = output.with_suffix(".json")
|
|
meta = {
|
|
"checkpoint": str(Path(checkpoint_path).resolve()),
|
|
"prompt": prompt,
|
|
"negative_prompt": negative_prompt,
|
|
"width": width,
|
|
"height": height,
|
|
"steps": steps,
|
|
"cfg": cfg,
|
|
"seed": seed,
|
|
"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}")
|