| 26 | |
| 27 | |
| 28 | class TrainerModuleLatentFM(LightningModule): |
| 29 | def __init__( |
| 30 | self, |
| 31 | # models |
| 32 | fm_cfg: dict, |
| 33 | noising_step: int = -1, |
| 34 | start_from_noise: bool = False, |
| 35 | # first stage |
| 36 | first_stage: dict = None, |
| 37 | # lora |
| 38 | lora_cfg: dict = None, |
| 39 | # conditioning |
| 40 | cond_stage_cfg: dict = None, |
| 41 | context_key: str = None, |
| 42 | cond_dropout: float = 0.0, |
| 43 | conditioning_key: str = None, |
| 44 | # training |
| 45 | lr: float = 1e-4, |
| 46 | weight_decay: float = 0., |
| 47 | sampling_steps: int = 50, |
| 48 | ema_rate: float = 0.999, |
| 49 | ema_update_every: int = 100, |
| 50 | ema_update_after_step: int = 1000, |
| 51 | use_ema_for_sampling: bool = True, |
| 52 | lr_scheduler_cfg: dict = None, |
| 53 | # logging |
| 54 | n_images_to_vis: int = 16, |
| 55 | log_grad_norm: bool = False, |
| 56 | metric_tracker_cfg: dict = None, |
| 57 | visualizer: dict = None, |
| 58 | ): |
| 59 | """ |
| 60 | Args: |
| 61 | fm_cfg: Flow matching model config. |
| 62 | noising_step: Forward diffusion noising step with linear schedule |
| 63 | of Ho et al. Set to -1 to disable. |
| 64 | start_from_noise: Whether to start from noise with low-res image as |
| 65 | conditioning (FM) or directly from low-res image (IC-FM). |
| 66 | first_stage: First stage config, if None, identity is used. |
| 67 | lora_cfg: LoRA config, if None, no LoRA is used. |
| 68 | cond_stage_cfg: Conditioning stage config, if None, no conditioning is used. |
| 69 | context_key: Context conditioning signal, concatenated to the input. |
| 70 | conditioning_key: Key in the batch to use for conditioning. |
| 71 | cond_dropout: Dropout rate for conditioning. |
| 72 | lr: Learning rate. |
| 73 | weight_decay: Weight decay. |
| 74 | sampling_steps: Number of sampling steps for inference. |
| 75 | ema_rate: EMA rate. |
| 76 | ema_update_every: EMA update rate (every n steps). |
| 77 | ema_update_after_step: EMA update start after n steps. |
| 78 | use_ema_for_sampling: Whether to use the EMA model for sampling. |
| 79 | lr_scheduler_cfg: Learning rate scheduler config. |
| 80 | n_images_to_vis: Number of images to visualize. |
| 81 | log_grad_norm: Whether to log gradient norm. |
| 82 | metric_tracker_cfg: Metric tracker config, if None, no metrics are tracked. |
| 83 | visualizer: Visualizer config, if None, no visualization is done. |
| 84 | """ |
| 85 | super().__init__() |