Add D-OPSD as a distillation handeling option for MiniMax H3 ref2va

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

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()