| |
| """Single-prompt SD3-M / SD3.5-M generation with the matched released adapter. |
| |
| Install the image package from https://github.com/TWWinde/ReSeed first. |
| No remote Python model code is loaded by this script. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--weights", type=Path, required=True, help="downloaded single-seed snapshot") |
| parser.add_argument("--model", choices=["sd3m", "sd35m"], required=True) |
| parser.add_argument("--base", type=Path, required=True, help="downloaded upstream teacher snapshot") |
| parser.add_argument("--prompt", required=True) |
| parser.add_argument("--seed", type=int, default=0, help="generation noise seed, not training seed") |
| parser.add_argument("--output", type=Path, required=True) |
| parser.add_argument("--without-reseed", action="store_true") |
| args = parser.parse_args() |
| if args.output.exists(): |
| raise FileExistsError("refusing to overwrite an existing image") |
| import torch |
| from diffusers import StableDiffusion3Pipeline, SD3Transformer2DModel |
| from reseed.models.sd3_residual_bias import load_sd3_residual_bias_adapter |
|
|
| entry = json.loads((args.weights/"release_manifest.json").read_text())["models"][args.model] |
| transformer = SD3Transformer2DModel.from_pretrained( |
| str(args.weights/entry["student_path"]/"student"), torch_dtype=torch.bfloat16) |
| pipe = StableDiffusion3Pipeline.from_pretrained(str(args.base), transformer=transformer, |
| torch_dtype=torch.bfloat16).to("cuda") |
| if abs(float(pipe.scheduler.config.shift)-3.) > 1e-7 or pipe.scheduler.config.get("use_dynamic_shifting", False): |
| raise ValueError("this release uses static FlowMatch shift=3, not dynamic shifting") |
| for component in pipe.components.values(): |
| if isinstance(component, torch.nn.Module): |
| component.eval().requires_grad_(False) |
| adapter = None |
| if not args.without_reseed: |
| adapter, _ = load_sd3_residual_bias_adapter(pipe.transformer, args.weights/entry["adapter_path"]) |
| adapter.to("cuda").eval().requires_grad_(False) |
| adapter.enable() |
| def callback(_pipe, index, _timestep, kwargs): |
| if adapter is not None: |
| adapter.disable() |
| return kwargs |
| shape = (1, int(transformer.config.in_channels), 1024//pipe.vae_scale_factor, 1024//pipe.vae_scale_factor) |
| noise = torch.randn(shape, generator=torch.Generator(device="cpu").manual_seed(args.seed), dtype=torch.float32) |
| try: |
| with torch.inference_mode(): |
| image = pipe(prompt=[args.prompt], negative_prompt=[""], height=1024, width=1024, |
| sigmas=[1., .75, .5, .25], guidance_scale=1., max_sequence_length=256, |
| latents=noise.to(device="cuda", dtype=torch.bfloat16), |
| callback_on_step_end=callback if adapter is not None else None, |
| callback_on_step_end_tensor_inputs=["latents"]).images[0] |
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| image.save(args.output) |
| finally: |
| if adapter is not None: |
| adapter.disable(); adapter.close() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
