File size: 2,426 Bytes
2ca14fe a30f2d4 2ca14fe a30f2d4 2ca14fe a30f2d4 2ca14fe | 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 | """
T.I.T.S.2 inference — rectified flow, ~30 Euler steps, 256x256.
Usage:
python sample2.py --checkpoint checkpoints2/latest.pt --prompt "a red bus on a city street" \
--num_images 4 --guidance 4.0 --out out.png
"""
import argparse
import torch
from safetensors import safe_open
from safetensors.torch import load_file
from torchvision.utils import save_image
from model2 import TITS2
from preview import decode, sample_latents
from text_encoder2 import FrozenCLIPTextEncoder
def parse_args():
p = argparse.ArgumentParser()
p.add_argument("--checkpoint", type=str, required=True)
p.add_argument("--prompt", type=str, required=True)
p.add_argument("--num_images", type=int, default=1)
p.add_argument("--steps", type=int, default=30)
p.add_argument("--guidance", type=float, default=4.0)
p.add_argument("--seed", type=int, default=None)
p.add_argument("--raw", action="store_true", help="use raw weights instead of EMA")
p.add_argument("--out", type=str, default="sample2.png")
p.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu")
return p.parse_args()
def main():
args = parse_args()
text_encoder = FrozenCLIPTextEncoder(device=args.device)
if args.checkpoint.endswith(".safetensors"):
with safe_open(args.checkpoint, framework="pt") as f:
meta = f.metadata() or {}
cfg = {k: int(meta.get(k, d)) for k, d in (("dim", 576), ("depth", 12), ("heads", 9))}
step = meta.get("step", "?")
state = {k: v.float() for k, v in load_file(args.checkpoint).items()}
else:
ck = torch.load(args.checkpoint, map_location=args.device, weights_only=False)
cfg = {k: ck[k] for k in ("dim", "depth", "heads")}
step = ck["step"]
state = ck["model_state_dict"] if args.raw else ck["ema_state_dict"]
model = TITS2(text_dim=text_encoder.embed_dim, **cfg).to(args.device)
model.load_state_dict(state)
model.eval()
lat = sample_latents(model, text_encoder, [args.prompt] * args.num_images,
steps=args.steps, guidance=args.guidance,
device=args.device, seed=args.seed)
imgs = decode(lat, args.device)
save_image(imgs, args.out, nrow=int(args.num_images**0.5) or 1)
print(f"saved {args.num_images} image(s) to {args.out} (step {step})")
if __name__ == "__main__":
main()
|