| """ |
| ArtiGen Sampling — Flow Matching ODE Solver. |
| Rectified flow allows efficient 1-4 step generation via Euler solver. |
| """ |
| import torch |
| import torch.nn.functional as F |
| import numpy as np |
|
|
| try: |
| from artigen.model import ArtiGen |
| except ImportError: |
| from model import ArtiGen |
|
|
|
|
| def euler_solver(model, z_noise, text_embed, num_steps=4, device='cpu', cfg_scale=1.0): |
| model.eval() |
| z = z_noise.clone() |
| dt = -1.0 / num_steps |
| with torch.no_grad(): |
| for i in range(num_steps): |
| t = torch.ones(z.shape[0], device=device) * (1.0 + i * dt) |
| if cfg_scale > 1.0: |
| v_cond, _ = model(z, t, text_embed, return_asdl=False) |
| v_uncond, _ = model(z, t, torch.zeros_like(text_embed), return_asdl=False) |
| v = v_uncond + cfg_scale * (v_cond - v_uncond) |
| else: |
| v, _ = model(z, t, text_embed, return_asdl=False) |
| z = z + dt * v |
| return z |
|
|
|
|
| def sample(model, text_embed, latent_shape=(4, 32, 32), num_steps=4, device='cpu', cfg_scale=1.5): |
| C, H, W = latent_shape |
| B = text_embed.shape[0] |
| z_noise = torch.randn(B, C, H, W, device=device) |
| z0 = euler_solver(model, z_noise, text_embed, num_steps=num_steps, device=device, cfg_scale=cfg_scale) |
| return z0 |
|
|
|
|
| def decode_with_vae(vae, z, output_type='pil'): |
| img = torch.randn(z.shape[0], 3, 256, 256) |
| return img |
|
|
