Spaces:
Running on Zero
Running on Zero
| # Copyright (c) 2025 Hansheng Chen | |
| import torch | |
| from abc import abstractmethod | |
| from copy import deepcopy | |
| from accelerate import init_empty_weights | |
| from mmgen.models.builder import build_module | |
| from .base import BaseModel | |
| from lakonlab.utils import clone_params, rgetattr, tie_untrained_submodules | |
| def train_fwd_bwd(model, args, kwargs, loss_scaler=None): | |
| is_multistep = rgetattr(model, 'is_multistep', False) | |
| if is_multistep: | |
| step_states, log_vars = model(*args, return_step_states=True, **kwargs) | |
| loss = 0 | |
| step_id = 0 | |
| while not step_states['terminate']: | |
| step_loss, step_log_vars, step_states = model( | |
| *args, return_loss=True, step_states=step_states, **kwargs) | |
| if step_states['detachable']: | |
| step_loss.backward() if loss_scaler is None else loss_scaler.scale(step_loss).backward() | |
| step_loss.detach_() | |
| loss = loss + step_loss | |
| for k, v in step_log_vars.items(): | |
| if k in log_vars: | |
| log_vars[k] += v | |
| else: | |
| log_vars[k] = v | |
| step_id += 1 | |
| else: | |
| loss, log_vars = model(*args, return_loss=True, **kwargs) | |
| if isinstance(loss, torch.Tensor) and loss.requires_grad: | |
| loss.backward() if loss_scaler is None else loss_scaler.scale(loss).backward() | |
| return log_vars | |
| class BaseDiffusion(BaseModel): | |
| """Base class providing the common training interface for diffusion models. Optionally supports: | |
| - Teacher model for distillation training | |
| - EMA version of the diffusion model | |
| - Multi-step diffusion training | |
| - Image/video patching for patch-wise GMFlow | |
| """ | |
| def __init__(self, | |
| diffusion=dict(type='GaussianFlow'), | |
| diffusion_use_ema=False, | |
| tie_ema=True, | |
| teacher=None, | |
| tie_teacher=False, | |
| patch_size=1, | |
| inference_only=False, | |
| train_cfg=None, | |
| test_cfg=None): | |
| super().__init__() | |
| # order matters: teacher must be built before diffusion for FSDP tying | |
| if teacher is not None and not inference_only: | |
| teacher.update(train_cfg=train_cfg, test_cfg=test_cfg) | |
| self.teacher = build_module(teacher) | |
| else: | |
| self.teacher = None | |
| diffusion.update(train_cfg=train_cfg, test_cfg=test_cfg) | |
| self.diffusion = build_module(diffusion) | |
| if self.teacher is not None and tie_teacher: | |
| tie_untrained_submodules(self.diffusion, self.teacher, tie_tgt_lora_base_layer=True) | |
| self.patch_size = patch_size | |
| self.diffusion_use_ema = diffusion_use_ema | |
| if self.diffusion_use_ema: | |
| if inference_only: | |
| self.diffusion_ema = self.diffusion | |
| else: | |
| diffusion_ema = deepcopy(diffusion) | |
| if isinstance(diffusion_ema.get('denoising', None), dict): | |
| diffusion_ema['denoising'].pop('pretrained', None) | |
| with init_empty_weights(): | |
| self.diffusion_ema = build_module(diffusion_ema) | |
| if tie_ema: | |
| tie_untrained_submodules(self.diffusion_ema, self.diffusion) | |
| clone_params(self.diffusion_ema, self.diffusion) | |
| self.train_cfg = dict() if train_cfg is None else deepcopy(train_cfg) | |
| self.test_cfg = dict() if test_cfg is None else deepcopy(test_cfg) | |
| def patchify(self, x): | |
| if isinstance(self.patch_size, int) and self.patch_size == 1: | |
| return x | |
| if x.dim() == 4: | |
| if isinstance(self.patch_size, int): | |
| ph = pw = self.patch_size | |
| else: | |
| assert len(self.patch_size) == 2 | |
| ph, pw = self.patch_size | |
| bs, c, h, w = x.size() | |
| x = x.reshape( | |
| bs, c, h // ph, ph, w // pw, pw | |
| ).permute( | |
| 0, 1, 3, 5, 2, 4 | |
| ).reshape( | |
| bs, c * ph * pw, h // ph, w // pw) | |
| elif x.dim() == 5: | |
| if isinstance(self.patch_size, int): | |
| pt = ph = pw = self.patch_size | |
| else: | |
| assert len(self.patch_size) == 3 | |
| pt, ph, pw = self.patch_size | |
| bs, c, t, h, w = x.size() | |
| x = x.reshape( | |
| bs, c, t // pt, pt, h // ph, ph, w // pw, pw | |
| ).permute( | |
| 0, 1, 3, 5, 7, 2, 4, 6 | |
| ).reshape( | |
| bs, c * pt * ph * pw, t // pt, h // ph, w // pw) | |
| else: | |
| raise ValueError(f'Unsupported input dimension {x.dim()}. Expected 4 or 5 dimensions.') | |
| return x | |
| def unpatchify(self, x): | |
| if isinstance(self.patch_size, int) and self.patch_size == 1: | |
| return x | |
| if x.dim() == 4: | |
| if isinstance(self.patch_size, int): | |
| ph = pw = self.patch_size | |
| else: | |
| assert len(self.patch_size) == 2 | |
| ph, pw = self.patch_size | |
| bs, c, h, w = x.size() | |
| x = x.reshape( | |
| bs, c // (ph * pw), ph, pw, h, w | |
| ).permute( | |
| 0, 1, 4, 2, 5, 3 | |
| ).reshape( | |
| bs, c // (ph * pw), h * ph, w * pw) | |
| elif x.dim() == 5: | |
| if isinstance(self.patch_size, int): | |
| pt = ph = pw = self.patch_size | |
| else: | |
| assert len(self.patch_size) == 3 | |
| pt, ph, pw = self.patch_size | |
| bs, c, t, h, w = x.size() | |
| x = x.reshape( | |
| bs, c // (pt * ph * pw), pt, ph, pw, t, h, w | |
| ).permute( | |
| 0, 1, 5, 2, 6, 3, 7, 4 | |
| ).reshape( | |
| bs, c // (pt * ph * pw), t * pt, h * ph, w * pw) | |
| else: | |
| raise ValueError(f'Unsupported input dimension {x.dim()}. Expected 4 or 5 dimensions.') | |
| return x | |
| def _prepare_train_minibatch_args(self, data, running_status=None): | |
| """ | |
| Prepare the arguments for the training minibatch. | |
| Args: | |
| data (dict): The input data for the training step. | |
| running_status (dict): The running status for the training step. | |
| Returns: | |
| tuple: A tuple containing the batch size, diffusion arguments, and diffusion keyword arguments. | |
| """ | |
| def train_minibatch(self, data, loss_scaler=None, running_status=None): | |
| bs, diffusion_args, diffusion_kwargs = self._prepare_train_minibatch_args(data, running_status) | |
| log_vars = train_fwd_bwd(self.diffusion, diffusion_args, diffusion_kwargs, loss_scaler) | |
| return log_vars, bs | |
| def val_step(self, data, test_cfg_override=dict(), **kwargs): | |
| """Perform a validation step. | |
| Args: | |
| data (dict): The input data for the validation step. | |
| test_cfg_override (dict): Override configuration for the test. | |
| Returns: | |
| dict: A dictionary containing the number of samples and predicted outputs. | |
| """ | |