Spaces:
Running on Zero
Running on Zero
| """End-to-end text-to-image generation (used by inference and train-time validation).""" | |
| from pathlib import Path | |
| import torch | |
| from iris3b.flow.solver import FlowDPMSolver | |
| from iris3b.models.dit import IrisDiT | |
| from iris3b.text.base import TextEncoder | |
| def load_for_inference(path: str | Path) -> tuple[dict, dict[str, torch.Tensor]]: | |
| """``(config dict, weights)`` from a training checkpoint or an exported directory. | |
| A training ``.pth`` contributes its embedded config and its EMA weights, | |
| or its raw weights when it was trained without EMA. The checkpoint is | |
| memory-mapped, so its optimizer state is never read. An exported directory | |
| holds ``config.yaml`` and ``model.safetensors``. | |
| """ | |
| path = Path(path) | |
| if path.is_dir(): | |
| from omegaconf import OmegaConf | |
| from safetensors.torch import load_file | |
| raw = OmegaConf.to_container(OmegaConf.load(path / "config.yaml")) | |
| return raw, load_file(path / "model.safetensors") | |
| payload = torch.load(path, map_location="cpu", mmap=True, weights_only=False) | |
| if "config" not in payload: | |
| raise ValueError(f"{path} has no embedded config; it is not a training checkpoint") | |
| key = "state_dict_ema" if "state_dict_ema" in payload else "state_dict" | |
| return payload["config"], payload[key] | |
| def generate( | |
| model: IrisDiT, | |
| text_encoder: TextEncoder, | |
| prompts: list[str], | |
| height: int, | |
| width: int, | |
| steps: int = 100, | |
| order: int = 2, | |
| cfg_scale: float = 3.0, | |
| cfg_interval: tuple[float, float] = (0.0, 1.0), | |
| shift: float = 4.0, | |
| negative_prompt: str = "", | |
| generator: torch.Generator | None = None, | |
| device: torch.device | str = "cuda", | |
| noise: torch.Tensor | None = None, | |
| num_train_timesteps: int = 1000, | |
| prediction: str = "v", | |
| ) -> torch.Tensor: | |
| """Sample [B, C, H, W] images, clamped to [-1, 1]. | |
| Integrator state stays in fp32; the network is evaluated in its own | |
| parameter dtype. The CFG unconditional is ``text_encoder.null`` of the | |
| negative prompt, which the encoder formats exactly like a positive prompt | |
| (an empty negative prompt is the dropout null used in training). | |
| """ | |
| model_dtype = next(model.parameters()).dtype | |
| encoding = text_encoder.encode(prompts) | |
| cond = encoding.embeddings.to(device=device, dtype=model_dtype) | |
| cond_mask = encoding.mask.to(device=device) | |
| uncond, cfg_mask = None, None | |
| if cfg_scale != 1.0: | |
| null = text_encoder.null(negative_prompt) | |
| uncond = null.embeddings.to(device=device, dtype=model_dtype).expand( | |
| len(prompts), *([-1] * (null.embeddings.ndim - 1)) | |
| ) | |
| null_mask = null.mask.to(device=device).expand(len(prompts), -1) | |
| # the solver's CFG batch is cat([uncond, cond]); the mask follows suit | |
| cfg_mask = torch.cat([null_mask, cond_mask], dim=0) | |
| if noise is None: | |
| noise = torch.randn( | |
| len(prompts), model.cfg.in_channels, height, width, generator=generator, device=device | |
| ) | |
| z = noise.to(device=device, dtype=torch.float32) | |
| def model_fn(x: torch.Tensor, t_model: torch.Tensor, y: torch.Tensor) -> torch.Tensor: | |
| # the solver either passes `cond` itself or the doubled CFG batch | |
| out = model(x.to(model_dtype), t_model, y, y_mask=cond_mask if y is cond else cfg_mask) | |
| return out.x.float() | |
| solver = FlowDPMSolver( | |
| model_fn, | |
| num_timesteps=num_train_timesteps, | |
| cfg_scale=cfg_scale, | |
| cfg_interval=tuple(cfg_interval), | |
| prediction=prediction, | |
| ) | |
| sample = solver.sample(z, cond, uncond, steps=steps, order=order, shift=shift) | |
| return sample.clamp(-1, 1) | |