|
""" |
|
Generate images from trained model |
|
""" |
|
|
|
import argparse |
|
import pickle |
|
|
|
import jax |
|
import jax.numpy as jnp |
|
import matplotlib.pyplot as plt |
|
import yaml |
|
from flax import nnx |
|
from jax.experimental import ode |
|
|
|
from model import DiT, DiTConfig |
|
|
|
|
|
def parse_args(): |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument( |
|
"--config", type=str, default="config.yaml", help="Path to config file" |
|
) |
|
parser.add_argument( |
|
"--ckpt", type=str, default=None, help="Path to checkpoint file" |
|
) |
|
parser.add_argument("--seed", type=int, default=0, help="Random seed") |
|
return parser.parse_args() |
|
|
|
|
|
def load_config(config_path): |
|
with open(config_path) as f: |
|
config = yaml.safe_load(f) |
|
return config |
|
|
|
|
|
@jax.jit |
|
def sample_images(graphdef, state, rng): |
|
flow = nnx.merge(graphdef, state) |
|
|
|
def flow_fn(y, t): |
|
o = flow(y, t[None]) |
|
return o |
|
|
|
x = jax.random.normal(rng, shape=(16, 64, 64, 3), dtype=jnp.float32) |
|
o = ode.odeint(flow_fn, x, jnp.linspace(0, 1, 1000)) |
|
o = jnp.clip(o[-1], 0, 1) |
|
return o |
|
|
|
|
|
def plot_new_images(graphdef, state, seed): |
|
images = sample_images(graphdef, state, nnx.Rngs(seed)()) |
|
|
|
plt.figure(figsize=(2, 2)) |
|
for i in range(16): |
|
plt.subplot(4, 4, i + 1) |
|
plt.imshow(images[i]) |
|
plt.axis("off") |
|
plt.subplots_adjust(left=0, bottom=0, top=1, right=1, wspace=0, hspace=0) |
|
plt.savefig(f"samples.png") |
|
plt.close() |
|
|
|
|
|
def main(): |
|
args = parse_args() |
|
config = load_config(args.config) |
|
|
|
dit_config = DiTConfig( |
|
input_dim=config["model"]["input_dim"], |
|
hidden_dim=config["model"]["hidden_dim"], |
|
num_blocks=config["model"]["num_blocks"], |
|
num_heads=config["model"]["num_heads"], |
|
patch_size=config["model"]["patch_size"], |
|
patch_stride=config["model"]["patch_stride"], |
|
time_freq_dim=config["model"]["time_freq_dim"], |
|
time_max_period=config["model"]["time_max_period"], |
|
mlp_ratio=config["model"]["mlp_ratio"], |
|
use_bias=config["model"]["use_bias"], |
|
padding=config["model"]["padding"], |
|
pos_embed_cls_token=config["model"]["pos_embed_cls_token"], |
|
pos_embed_extra_tokens=config["model"]["pos_embed_extra_tokens"], |
|
) |
|
|
|
abstract_flow = nnx.eval_shape(lambda: DiT(dit_config, rngs=nnx.Rngs(0))) |
|
graphdef, _ = nnx.split(abstract_flow) |
|
with open(args.ckpt, "rb") as f: |
|
state = pickle.load(f, fix_imports=True) |
|
if "time_embedding" not in state: |
|
state = state[0] |
|
plot_new_images(graphdef, state, args.seed) |
|
|
|
|
|
if __name__ == "__main__": |
|
main() |
|
|
