File size: 3,422 Bytes
51133c2 | 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 | #!/usr/bin/env python3
"""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() # Step zero just finished; never activate on steps 1–3.
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()
|