Spaces:
Running on Zero
Running on Zero
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: OpenMDW-1.1 | |
| """Fixed class-conditioned image sampling callback for DiT training.""" | |
| from __future__ import annotations | |
| from contextlib import nullcontext | |
| from typing import Any | |
| import torch | |
| import torchvision | |
| import wandb | |
| from cosmos_framework.callbacks.every_n import EveryN | |
| from cosmos_framework.model._base import ImaginaireModel | |
| from cosmos_framework.trainer import ImaginaireTrainer | |
| from cosmos_framework.utils import distributed | |
| class DiTImageSampleCallback(EveryN): | |
| """Generate fixed ImageNet class samples through ``model.generate_image``.""" | |
| def __init__( | |
| self, | |
| every_n: int = 5000, | |
| class_ids: list[int] | None = None, | |
| cfg_scales: list[float] | None = None, | |
| num_steps: int = 50, | |
| seed: int = 0, | |
| is_ema: bool = True, | |
| run_at_start: bool = False, | |
| ) -> None: | |
| super().__init__(every_n=every_n, run_at_start=run_at_start) | |
| self.class_ids = class_ids or [0, 1, 2, 3] | |
| self.cfg_scales = cfg_scales or [1.0, 1.25, 1.5, 2.0] | |
| self.num_steps = num_steps | |
| self.seed = seed | |
| self.is_ema = is_ema | |
| self.rank = distributed.get_rank() | |
| def every_n_impl( | |
| self, | |
| trainer: ImaginaireTrainer, | |
| model: ImaginaireModel, | |
| data_batch: dict[str, torch.Tensor], | |
| output_batch: dict[str, torch.Tensor], | |
| loss: torch.Tensor, | |
| iteration: int, | |
| ) -> None: | |
| del trainer, data_batch, output_batch, loss | |
| if not hasattr(model, "generate_image"): | |
| raise AttributeError("DiTImageSampleCallback requires model.generate_image().") | |
| if self.is_ema and not model.config.ema.enabled: | |
| return | |
| was_training = model.training | |
| context: Any = model.ema_scope("dit_image_sample") if self.is_ema else nullcontext() | |
| generated_rows: list[torch.Tensor] = [] | |
| seed_list = [self.seed + sample_idx for sample_idx in range(len(self.class_ids))] | |
| try: | |
| with context: | |
| for cfg_scale in self.cfg_scales: | |
| images = model.generate_image( | |
| class_ids=self.class_ids, | |
| num_steps=self.num_steps, | |
| cfg_scale=cfg_scale, | |
| seed=seed_list, | |
| ) # [B,3,H,W] | |
| if self.rank == 0: | |
| generated_rows.append(images.detach().float().cpu()) # [B,3,H,W] | |
| finally: | |
| if was_training: | |
| model.train() | |
| if self.rank != 0 or wandb.run is None or not generated_rows: | |
| return | |
| grid_images = torch.cat(generated_rows, dim=0) # [R*B,3,H,W] | |
| grid = torchvision.utils.make_grid(grid_images, nrow=len(self.class_ids), padding=2, normalize=False) # [3,H,W] | |
| grid_np = grid.clamp(0.0, 1.0).permute(1, 2, 0).numpy() # [H,W,3] | |
| tag = "ema" if self.is_ema else "reg" | |
| caption = f"classes={self.class_ids}, cfg={self.cfg_scales}, steps={self.num_steps}, seed={self.seed}, {tag}" | |
| wandb.log( | |
| {f"dit_image_sample/{tag}": wandb.Image(grid_np, caption=caption)}, | |
| step=iteration, | |
| ) | |