From da79ebce991c16bc6d343348085d0f7173e367ab Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Wed, 26 Aug 2026 11:06:33 -0600 Subject: [PATCH] Add D-OPSD as a distillation handeling option for MiniMax H3 ref2va --- .../diffusion_models/minimax_h3/minimax_h3.py | 58 +++++++++ extensions_built_in/sd_trainer/SDTrainer.py | 21 ++++ jobs/process/BaseSDTrainProcess.py | 16 +++ toolkit/data_loader.py | 1 + toolkit/data_transfer_object/data_loader.py | 24 ++++ toolkit/dataloader_mixins.py | 111 +++++++++++++++++- toolkit/models/base_model.py | 4 + toolkit/stable_diffusion_model.py | 4 + ui/src/app/jobs/new/options.tsx | 52 ++++++-- version.py | 2 +- 10 files changed, 278 insertions(+), 15 deletions(-) diff --git a/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py b/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py index 358b0e2..1ec8619 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py +++ b/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py @@ -1297,6 +1297,12 @@ class MinimaxH3Ref2VAModel(MinimaxH3Model): # control VIDEOS are cached like dataset items and consumed as # multi-frame reference blocks self.supports_video_control_images = True + # D-OPSD: a no-grad teacher pass with the target as its own reference + # becomes the training target for the reference-free student pass + self.dopsd = bool(self.model_config.model_kwargs.get("dopsd", False)) + if self.dopsd: + self.dopsd_self_ref = True + self.require_pixel_tensor_cache = True def _dit_component(self) -> str: partition = str( @@ -1319,6 +1325,10 @@ class MinimaxH3Ref2VAModel(MinimaxH3Model): # pixel area with its own aspect kept, then encoded on its own grid if batch is None: return None, None, (), () + if getattr(batch, "dopsd_teacher_pass", False): + return self._build_dopsd_teacher_condition( + batch, latent_shape, device, dtype + ) controls_per_item = None if batch.control_tensor_list is not None: controls_per_item = batch.control_tensor_list @@ -1400,6 +1410,54 @@ class MinimaxH3Ref2VAModel(MinimaxH3Model): cond_audio = torch.cat(audio_rows, dim=1) if audio_rows else None return torch.cat(all_rows, dim=1), cond_audio, (), tuple(blocks) + def _build_dopsd_teacher_condition(self, batch, latent_shape, device, dtype): + """D-OPSD teacher: the target is its own (only) reference, built from + the batch's cached pixel tensors; its soundtrack rides as clean + reference audio rows.""" + if batch.tensor is None: + raise ValueError( + "D-OPSD teacher pass needs pixel tensors on the batch; set " + "cache_tensors_to_disk: true on this dataset" + ) + px = batch.tensor.detach().to(device, torch.float32) + if px.ndim == 4: + # image items: single-frame ref, or a static clip for image_refs_as_video + frames = px.unsqueeze(2) + as_video_frames = self._image_ref_video_frames() + if as_video_frames: + frames = frames.expand(-1, -1, as_video_frames, -1, -1).contiguous() + else: + # video items: (B, T, C, H, W) -> (B, C, T, H, W) + frames = px.permute(0, 2, 1, 3, 4).contiguous() + ref_latents = self.encode_keyframe_latents(frames) + ref_noise = torch.randn_like(ref_latents) + ref_latents = ( + KEYFRAME_NOISE_AUG_T * ref_latents + + (1.0 - KEYFRAME_NOISE_AUG_T) * ref_noise + ) + audio_rows = None + a_lat = 0 + if ( + batch.dataset_config is not None + and batch.dataset_config.do_audio + and batch.audio_latents is not None + and getattr(batch, "num_frames", 1) > 1 + ): + a_lat = packing.audio_latent_num_frames(batch.num_frames) + audio_rows = torch.stack( + [ + self._fit_audio_rows( + batch.audio_latents[b].detach().to(device, torch.float32), + a_lat, + ) + for b in range(batch.audio_latents.shape[0]) + ] + ).to(dtype) + blocks = ( + (ref_latents.shape[2], ref_latents.shape[3], ref_latents.shape[4], a_lat), + ) + return patchify_video_latents(ref_latents).to(dtype), audio_rows, (), blocks + @torch.no_grad() def _encode_ref_video_for_sampling(self, path: str, gen_config) -> torch.Tensor: """Decode a reference video with the SAME temporal treatment training diff --git a/extensions_built_in/sd_trainer/SDTrainer.py b/extensions_built_in/sd_trainer/SDTrainer.py index 537c3bc..5fba071 100644 --- a/extensions_built_in/sd_trainer/SDTrainer.py +++ b/extensions_built_in/sd_trainer/SDTrainer.py @@ -309,6 +309,9 @@ class SDTrainer(BaseSDTrainProcess): if self.train_config.do_prior_divergence: self.do_prior_prediction = True + if getattr(self.sd, 'dopsd_self_ref', False): + # D-OPSD: the teacher (prior) prediction is the training target + self.do_prior_prediction = True # move vae to device if we did not cache latents if not self.is_latents_cached: self.sd.vae.eval() @@ -1995,6 +1998,19 @@ class SDTrainer(BaseSDTrainProcess): [blank_embeds] * noisy_latents.shape[0] ) + is_dopsd = getattr(self.sd, 'dopsd_self_ref', False) + if is_dopsd: + # teacher embeds: caption + vision block naming the item as its own reference + if batch.dopsd_prompt_embeds is None: + raise ValueError( + "D-OPSD requires cached text embeddings; enable " + "train.cache_text_embeddings" + ) + prior_embeds_to_use = batch.dopsd_prompt_embeds.clone().detach().to( + self.device_torch, dtype=dtype + ) + batch.dopsd_teacher_pass = True + # joint audio models stash their audio pred on the batch. # Give this pass its own slot so the preservation loss # can pair it with the preservation pass below. @@ -2012,6 +2028,11 @@ class SDTrainer(BaseSDTrainProcess): conditioned_prompts=conditioned_prompts ) batch.audio_pred_slot = None + if is_dopsd: + batch.dopsd_teacher_pass = False + if batch.audio_pred_prior is not None: + # the teacher's audio prediction is the audio target too + batch.audio_target = batch.audio_pred_prior.detach() if prior_pred is not None: prior_pred = prior_pred.detach() diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py index d7de108..28dfa42 100644 --- a/jobs/process/BaseSDTrainProcess.py +++ b/jobs/process/BaseSDTrainProcess.py @@ -2247,6 +2247,22 @@ class BaseSDTrainProcess(BaseTrainProcess): ### HOOk ### self.before_dataset_load() + if getattr(self.sd, 'require_pixel_tensor_cache', False): + # model needs pixel tensors at train time even with cached latents; + # storing them changes the latent cache key (first run re-caches) + for ds_list in [self.datasets, self.datasets_reg]: + if ds_list is None: + continue + for ds in ds_list: + if not (ds.cache_latents or ds.cache_latents_to_disk): + # live-loading datasets already have pixels on the batch + continue + if not ds.cache_tensors_to_disk: + print_acc( + f"Model requires cached pixel tensors: forcing " + f"cache_tensors_to_disk on dataset {ds.folder_path}" + ) + ds.cache_tensors_to_disk = True # load datasets if passed in the root process if self.datasets is not None: self.data_loader = get_dataloader_from_datasets(self.datasets, self.train_config.batch_size, self.sd) diff --git a/toolkit/data_loader.py b/toolkit/data_loader.py index 7be76be..38b7ad4 100644 --- a/toolkit/data_loader.py +++ b/toolkit/data_loader.py @@ -539,6 +539,7 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin dataset_root=dataset_folder, encode_control_in_text_embeddings=self.sd.encode_control_in_text_embeddings if self.sd else False, encode_first_frame_in_text_embeddings=getattr(self.sd, 'encode_first_frame_in_text_embeddings', False) if self.sd else False, + dopsd_self_ref=getattr(self.sd, 'dopsd_self_ref', False) if self.sd else False, text_embedding_space_version=self.sd.text_embedding_space_version if self.sd else "sd1", te_padding_side=self.sd.te_padding_side if self.sd else "right", latent_space_version=latent_space_version, diff --git a/toolkit/data_transfer_object/data_loader.py b/toolkit/data_transfer_object/data_loader.py index d4051f3..8c7dfc6 100644 --- a/toolkit/data_transfer_object/data_loader.py +++ b/toolkit/data_transfer_object/data_loader.py @@ -83,6 +83,8 @@ class FileItemDTO( self.encode_first_frame_in_text_embeddings = kwargs.get( "encode_first_frame_in_text_embeddings", False ) + # D-OPSD: also cache teacher embeds with the item's own media as reference 1 + self.dopsd_self_ref = kwargs.get("dopsd_self_ref", False) self.te_padding_side = kwargs.get("te_padding_side", "right") self.latent_space_version = kwargs.get("latent_space_version", "sd1") self.text_embedding_space_version = kwargs.get("text_embedding_space_version", "sd1") @@ -265,6 +267,8 @@ class DataLoaderBatchDTO: # noisy/sigma bookkeeping) directly. The trainer sets this around # its prior / guidance-unconditional / preservation passes. self.audio_pred_slot: Union[str, None] = None + # set by the trainer around the D-OPSD teacher pass + self.dopsd_teacher_pass: bool = False # noisy audio rows and audio sigma of the primary pass, used to # rebuild the clean audio estimate for perceptual losses self.audio_noisy: Union[torch.Tensor, None] = None @@ -325,6 +329,8 @@ class DataLoaderBatchDTO: self.prompt_embeds: Union[PromptEmbeds, None] = None # diff output preservation embeds (trigger word replaced with class) self.dop_prompt_embeds: Union[PromptEmbeds, None] = None + # D-OPSD teacher embeds (trigger word replaced with the self-reference token) + self.dopsd_prompt_embeds: Union[PromptEmbeds, None] = None # if self.file_items[0].control_tensor is not None: # if any have a control tensor, we concatenate them if any([x.control_tensor is not None for x in self.file_items]): @@ -517,6 +523,24 @@ class DataLoaderBatchDTO: self.dop_prompt_embeds = concat_prompt_embeds(dop_prompt_embeds_list, padding_side=padding_side) + if any([getattr(x, 'dopsd_prompt_embeds', None) is not None for x in self.file_items]): + # find one to use as a base + base_dopsd_prompt_embeds = None + for x in self.file_items: + if x.dopsd_prompt_embeds is not None: + base_dopsd_prompt_embeds = x.dopsd_prompt_embeds + break + dopsd_prompt_embeds_list = [] + for x in self.file_items: + if x.dopsd_prompt_embeds is None: + y = base_dopsd_prompt_embeds + else: + y = x.dopsd_prompt_embeds + dopsd_prompt_embeds_list.append(y) + padding_side = self.file_items[0].te_padding_side + + self.dopsd_prompt_embeds = concat_prompt_embeds(dopsd_prompt_embeds_list, padding_side=padding_side) + if any([x.audio_tensor is not None for x in self.file_items]): # find one to use as a base base_audio_tensor = None diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py index 59e592e..60e0a9e 100644 --- a/toolkit/dataloader_mixins.py +++ b/toolkit/dataloader_mixins.py @@ -324,6 +324,8 @@ class CaptionProcessingDTOMixin: self.caption_short: str = None # caption with the trigger word replaced by the diff output preservation class self.caption_dop: str = None + # D-OPSD teacher caption (trigger word replaced by /