File size: 6,284 Bytes
a5318e5 | 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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | from __future__ import annotations
import argparse
import hashlib
import json
import os
from pathlib import Path
import torch
from PIL import Image
from transformers import AutoTokenizer, T5EncoderModel
from solpix import AutoencoderDCSol, SolPixConfig, SolPixTransformer2D
TEXT_MODEL = "google/flan-t5-base"
PROMPTS = [
"Three Black men sharing french fries at a neighborhood diner, candid documentary photography.",
"A red fox standing in fresh snow beneath pine trees at winter dawn, wildlife photography.",
"A glass greenhouse filled with ferns after rain, soft natural light, botanical photograph.",
"A handmade cobalt blue teapot on a pale stone table, clean studio product photograph.",
"A white sailboat crossing a calm blue bay at golden hour, fine art landscape photograph.",
"An orange cat curled on a wooden chair in a sunlit bookshop, cozy editorial photograph.",
"A small street cafe reflected in wet pavement at night, warm window light, city photograph.",
"A wooden lighthouse on a rocky coast under a cloudy sky, atmospheric landscape photograph.",
"A bowl of ripe peaches on a kitchen counter, morning light, natural still life photograph.",
"A snow-covered cabin among tall pine trees at blue hour, quiet winter landscape photograph.",
"A baker placing fresh bread on a cooling rack in a bright kitchen, documentary photograph.",
"A goldfinch perched on a thin branch among spring blossoms, close-up wildlife photograph.",
"A red bicycle leaning against a brick wall on a leafy neighborhood street, lifestyle photograph.",
"A lemon cake with a slice cut out on a ceramic plate, bright tabletop food photograph.",
"A small observatory beneath a clear star-filled sky, distant mountains, night landscape photograph.",
]
def encode(text: str, tokenizer, encoder, device: torch.device):
tokens = tokenizer(text, return_tensors="pt", truncation=True, max_length=96)
ids = tokens["input_ids"].to(device)
mask = tokens["attention_mask"].to(device).bool()
with torch.inference_mode():
embeddings = encoder(input_ids=ids, attention_mask=mask).last_hidden_state
return embeddings.float(), mask
def main() -> None:
parser = argparse.ArgumentParser(description="Create individual SolPix release samples.")
parser.add_argument("--checkpoint", required=True)
parser.add_argument("--output-dir", required=True)
parser.add_argument("--seed", type=int, default=260926)
parser.add_argument("--steps", type=int, default=40)
parser.add_argument("--guidance-scale", type=float, default=3.5)
args = parser.parse_args()
root = Path(__file__).resolve().parent
os.environ.setdefault("HF_HOME", str(root / "hf_cache"))
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.bfloat16 if device.type == "cuda" and torch.cuda.is_bf16_supported() else torch.float32
checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
model = SolPixTransformer2D(SolPixConfig(**checkpoint["model_config"])).to(device).eval()
model_keys = model.state_dict()
weights = checkpoint.get("ema") or checkpoint["model"]
weights = {name: value for name, value in weights.items() if name in model_keys}
missing = set(model_keys) - set(weights)
if missing:
raise ValueError(f"Checkpoint is missing model weights: {sorted(missing)[:8]}")
model.load_state_dict(weights, strict=True)
tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL)
text_encoder = T5EncoderModel.from_pretrained(TEXT_MODEL, torch_dtype=dtype).to(device).eval()
text_encoder.requires_grad_(False)
empty, empty_mask = encode("", tokenizer, text_encoder, device)
vae = AutoencoderDCSol.from_pretrained(torch_dtype=dtype).to(device).eval()
scale = float(vae.config.scaling_factor or 0.41407)
time_grid = torch.linspace(1.0, 0.0, args.steps + 1, device=device)
records = []
for index, prompt in enumerate(PROMPTS, start=1):
conditional, conditional_mask = encode(prompt, tokenizer, text_encoder, device)
sample_seed = args.seed + index - 1
generator = torch.Generator(device="cpu").manual_seed(sample_seed)
latents = torch.randn(1, 32, 16, 16, generator=generator).to(device)
with torch.inference_mode():
for current, following in zip(time_grid[:-1], time_grid[1:]):
time = current.expand(1)
with torch.autocast(device_type=device.type, dtype=dtype, enabled=device.type == "cuda"):
v_cond = model(latents, time, conditional, conditional_mask)
v_uncond = model(latents, time, empty, empty_mask)
velocity = v_uncond + args.guidance_scale * (v_cond - v_uncond)
latents = latents + (following - current) * velocity.float()
decoded = vae.decode((latents / scale).to(dtype=dtype), return_dict=False)[0]
pixels = (decoded.float().squeeze(0).permute(1, 2, 0).cpu().clamp(-1, 1) * 127.5 + 127.5)
image = Image.fromarray(pixels.to(torch.uint8).numpy(), mode="RGB")
image_name = f"sample_{index:02d}.png"
image_path = output_dir / image_name
image.save(image_path, format="PNG", optimize=True)
record = {
"index": index,
"image": image_name,
"checkpoint_step": int(checkpoint.get("step", -1)),
"prompt": prompt,
"seed": sample_seed,
"sampling_steps": args.steps,
"guidance_scale": args.guidance_scale,
"text_encoder": TEXT_MODEL,
"decoder": AutoencoderDCSol.model_id,
"decoder_revision": AutoencoderDCSol.revision,
"sha256": hashlib.sha256(image_path.read_bytes()).hexdigest(),
}
(output_dir / f"sample_{index:02d}.json").write_text(json.dumps(record, indent=2) + "\n")
records.append(record)
print(f"[{index}/{len(PROMPTS)}] checkpoint={record['checkpoint_step']} {image_path}", flush=True)
(output_dir / "manifest.json").write_text(json.dumps(records, indent=2) + "\n")
if __name__ == "__main__":
main()
|