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

This commit is contained in:
Jaret Burkett
2026-08-26 11:06:33 -06:00
parent 8a912564ce
commit da79ebce99
10 changed files with 278 additions and 15 deletions

View File

@@ -1297,6 +1297,12 @@ class MinimaxH3Ref2VAModel(MinimaxH3Model):
# control VIDEOS are cached like dataset items and consumed as # control VIDEOS are cached like dataset items and consumed as
# multi-frame reference blocks # multi-frame reference blocks
self.supports_video_control_images = True 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: def _dit_component(self) -> str:
partition = str( partition = str(
@@ -1319,6 +1325,10 @@ class MinimaxH3Ref2VAModel(MinimaxH3Model):
# pixel area with its own aspect kept, then encoded on its own grid # pixel area with its own aspect kept, then encoded on its own grid
if batch is None: if batch is None:
return None, 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 controls_per_item = None
if batch.control_tensor_list is not None: if batch.control_tensor_list is not None:
controls_per_item = batch.control_tensor_list 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 cond_audio = torch.cat(audio_rows, dim=1) if audio_rows else None
return torch.cat(all_rows, dim=1), cond_audio, (), tuple(blocks) 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() @torch.no_grad()
def _encode_ref_video_for_sampling(self, path: str, gen_config) -> torch.Tensor: def _encode_ref_video_for_sampling(self, path: str, gen_config) -> torch.Tensor:
"""Decode a reference video with the SAME temporal treatment training """Decode a reference video with the SAME temporal treatment training

View File

@@ -309,6 +309,9 @@ class SDTrainer(BaseSDTrainProcess):
if self.train_config.do_prior_divergence: if self.train_config.do_prior_divergence:
self.do_prior_prediction = True 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 # move vae to device if we did not cache latents
if not self.is_latents_cached: if not self.is_latents_cached:
self.sd.vae.eval() self.sd.vae.eval()
@@ -1995,6 +1998,19 @@ class SDTrainer(BaseSDTrainProcess):
[blank_embeds] * noisy_latents.shape[0] [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. # joint audio models stash their audio pred on the batch.
# Give this pass its own slot so the preservation loss # Give this pass its own slot so the preservation loss
# can pair it with the preservation pass below. # can pair it with the preservation pass below.
@@ -2012,6 +2028,11 @@ class SDTrainer(BaseSDTrainProcess):
conditioned_prompts=conditioned_prompts conditioned_prompts=conditioned_prompts
) )
batch.audio_pred_slot = None 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: if prior_pred is not None:
prior_pred = prior_pred.detach() prior_pred = prior_pred.detach()

View File

@@ -2247,6 +2247,22 @@ class BaseSDTrainProcess(BaseTrainProcess):
### HOOk ### ### HOOk ###
self.before_dataset_load() 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 # load datasets if passed in the root process
if self.datasets is not None: if self.datasets is not None:
self.data_loader = get_dataloader_from_datasets(self.datasets, self.train_config.batch_size, self.sd) self.data_loader = get_dataloader_from_datasets(self.datasets, self.train_config.batch_size, self.sd)

View File

@@ -539,6 +539,7 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
dataset_root=dataset_folder, dataset_root=dataset_folder,
encode_control_in_text_embeddings=self.sd.encode_control_in_text_embeddings if self.sd else False, 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, 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", 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", te_padding_side=self.sd.te_padding_side if self.sd else "right",
latent_space_version=latent_space_version, latent_space_version=latent_space_version,

View File

@@ -83,6 +83,8 @@ class FileItemDTO(
self.encode_first_frame_in_text_embeddings = kwargs.get( self.encode_first_frame_in_text_embeddings = kwargs.get(
"encode_first_frame_in_text_embeddings", False "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.te_padding_side = kwargs.get("te_padding_side", "right")
self.latent_space_version = kwargs.get("latent_space_version", "sd1") self.latent_space_version = kwargs.get("latent_space_version", "sd1")
self.text_embedding_space_version = kwargs.get("text_embedding_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 # noisy/sigma bookkeeping) directly. The trainer sets this around
# its prior / guidance-unconditional / preservation passes. # its prior / guidance-unconditional / preservation passes.
self.audio_pred_slot: Union[str, None] = None 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 # noisy audio rows and audio sigma of the primary pass, used to
# rebuild the clean audio estimate for perceptual losses # rebuild the clean audio estimate for perceptual losses
self.audio_noisy: Union[torch.Tensor, None] = None self.audio_noisy: Union[torch.Tensor, None] = None
@@ -325,6 +329,8 @@ class DataLoaderBatchDTO:
self.prompt_embeds: Union[PromptEmbeds, None] = None self.prompt_embeds: Union[PromptEmbeds, None] = None
# diff output preservation embeds (trigger word replaced with class) # diff output preservation embeds (trigger word replaced with class)
self.dop_prompt_embeds: Union[PromptEmbeds, None] = None 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 self.file_items[0].control_tensor is not None:
# if any have a control tensor, we concatenate them # if any have a control tensor, we concatenate them
if any([x.control_tensor is not None for x in self.file_items]): 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) 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]): if any([x.audio_tensor is not None for x in self.file_items]):
# find one to use as a base # find one to use as a base
base_audio_tensor = None base_audio_tensor = None

View File

@@ -324,6 +324,8 @@ class CaptionProcessingDTOMixin:
self.caption_short: str = None self.caption_short: str = None
# caption with the trigger word replaced by the diff output preservation class # caption with the trigger word replaced by the diff output preservation class
self.caption_dop: str = None self.caption_dop: str = None
# D-OPSD teacher caption (trigger word replaced by <Picture 1>/<Video 1>)
self.caption_dopsd: str = None
dataset_config: DatasetConfig = kwargs.get('dataset_config', None) dataset_config: DatasetConfig = kwargs.get('dataset_config', None)
self.extra_values: List[float] = dataset_config.extra_values self.extra_values: List[float] = dataset_config.extra_values
@@ -378,6 +380,19 @@ class CaptionProcessingDTOMixin:
self.caption_dop = self.caption.replace( self.caption_dop = self.caption.replace(
self.trigger_word, self.dataset_config.diff_output_preservation_class self.trigger_word, self.dataset_config.diff_output_preservation_class
) )
if getattr(self, 'dopsd_self_ref', False):
# trigger word -> the self-reference token, or the token prepended
# when there is no trigger word
if self.trigger_word is not None:
self.caption_dopsd = self.caption.replace(
self.trigger_word, self.get_dopsd_ref_token()
)
else:
self.caption_dopsd = f"{self.get_dopsd_ref_token()} {self.caption}".strip()
def get_dopsd_ref_token(self: 'FileItemDTO') -> str:
# the item is always the only reference in D-OPSD mode
return "<Video 1>" if self.is_video else "<Picture 1>"
def get_caption( def get_caption(
self: 'FileItemDTO', self: 'FileItemDTO',
@@ -2111,13 +2126,18 @@ class TextEmbeddingFileItemDTOMixin:
self._blank_text_embedding_path: Union[str, None] = None self._blank_text_embedding_path: Union[str, None] = None
# DOP embeds for dropout steps (dropout caption with trigger replaced by class) # DOP embeds for dropout steps (dropout caption with trigger replaced by class)
self._dop_blank_text_embedding_path: Union[str, None] = None self._dop_blank_text_embedding_path: Union[str, None] = None
# D-OPSD teacher embeds (caption with trigger replaced by the self-reference
# token, encoded WITH the item's own image/video as the vision reference)
self.dopsd_prompt_embeds: Union[PromptEmbeds, None] = None
self._dopsd_text_embedding_path: Union[str, None] = None
self._dopsd_blank_text_embedding_path: Union[str, None] = None
self._loaded_text_embedding_path: Union[str, None] = None self._loaded_text_embedding_path: Union[str, None] = None
self._caption_was_dropped = False self._caption_was_dropped = False
self.is_text_embedding_cached = False self.is_text_embedding_cached = False
self.text_embedding_load_device = 'cpu' self.text_embedding_load_device = 'cpu'
self.text_embedding_version = 1 self.text_embedding_version = 1
def get_text_embedding_info_dict(self: 'FileItemDTO', caption_override=None, text_only=False): def get_text_embedding_info_dict(self: 'FileItemDTO', caption_override=None, text_only=False, dopsd_self_ref=False):
# make sure the caption is loaded here # make sure the caption is loaded here
# TODO: we need a way to cache all the other features like trigger words, DOP, etc. For now, we need to throw an error if not compatible. # TODO: we need a way to cache all the other features like trigger words, DOP, etc. For now, we need to throw an error if not compatible.
if self.caption is None: if self.caption is None:
@@ -2127,6 +2147,10 @@ class TextEmbeddingFileItemDTOMixin:
("text_embedding_space_version", self.text_embedding_space_version), ("text_embedding_space_version", self.text_embedding_space_version),
("text_embedding_version", self.text_embedding_version), ("text_embedding_version", self.text_embedding_version),
]) ])
if dopsd_self_ref:
# teacher embeds carry the item's own media as the vision reference
item["dopsd_self_ref"] = True
return item
# dropout embeds are encoded as plain text, keep control conditioning # dropout embeds are encoded as plain text, keep control conditioning
# out of their cache key # out of their cache key
if text_only: if text_only:
@@ -2150,11 +2174,11 @@ class TextEmbeddingFileItemDTOMixin:
item["first_frame_in_te"] = True item["first_frame_in_te"] = True
return item return item
def _build_text_embedding_path(self: 'FileItemDTO', caption_override=None, text_only=False): def _build_text_embedding_path(self: 'FileItemDTO', caption_override=None, text_only=False, dopsd_self_ref=False):
# we store text embeddings in a folder in same path as image called _text_embedding_cache # we store text embeddings in a folder in same path as image called _text_embedding_cache
img_dir = os.path.dirname(self.path) img_dir = os.path.dirname(self.path)
te_dir = os.path.join(img_dir, '_t_e_cache') te_dir = os.path.join(img_dir, '_t_e_cache')
hash_dict = self.get_text_embedding_info_dict(caption_override=caption_override, text_only=text_only) hash_dict = self.get_text_embedding_info_dict(caption_override=caption_override, text_only=text_only, dopsd_self_ref=dopsd_self_ref)
filename_no_ext = os.path.splitext(os.path.basename(self.path))[0] filename_no_ext = os.path.splitext(os.path.basename(self.path))[0]
# get base64 hash of md5 checksum of hash_dict # get base64 hash of md5 checksum of hash_dict
hash_input = json.dumps(hash_dict, sort_keys=True).encode('utf-8') hash_input = json.dumps(hash_dict, sort_keys=True).encode('utf-8')
@@ -2215,6 +2239,35 @@ class TextEmbeddingFileItemDTOMixin:
return self._dop_blank_text_embedding_path return self._dop_blank_text_embedding_path
def get_dopsd_text_embedding_path(self: 'FileItemDTO', recalculate=False):
if self._dopsd_text_embedding_path is not None and not recalculate:
return self._dopsd_text_embedding_path
# make sure the caption is loaded so caption_dopsd is built
if self.caption is None:
self.load_caption()
self._dopsd_text_embedding_path = self._build_text_embedding_path(
caption_override=self.caption_dopsd, dopsd_self_ref=True
)
return self._dopsd_text_embedding_path
def get_dopsd_dropout_caption(self: 'FileItemDTO'):
# dropout caption with the trigger word swapped for the self-reference token
dropout_caption = self.get_dropout_caption()
if self.trigger_word is not None:
return dropout_caption.replace(
self.trigger_word, self.get_dopsd_ref_token()
)
return f"{self.get_dopsd_ref_token()} {dropout_caption}".strip()
def get_dopsd_blank_text_embedding_path(self: 'FileItemDTO', recalculate=False):
if self._dopsd_blank_text_embedding_path is not None and not recalculate:
return self._dopsd_blank_text_embedding_path
# unlike DOP, dropout embeds keep the vision reference
self._dopsd_blank_text_embedding_path = self._build_text_embedding_path(
caption_override=self.get_dopsd_dropout_caption(), dopsd_self_ref=True
)
return self._dopsd_blank_text_embedding_path
def get_blank_text_embedding_path(self: 'FileItemDTO', recalculate=False): def get_blank_text_embedding_path(self: 'FileItemDTO', recalculate=False):
if self._blank_text_embedding_path is not None and not recalculate: if self._blank_text_embedding_path is not None and not recalculate:
return self._blank_text_embedding_path return self._blank_text_embedding_path
@@ -2235,6 +2288,8 @@ class TextEmbeddingFileItemDTOMixin:
self.prompt_embeds = None self.prompt_embeds = None
if self.dop_prompt_embeds is not None: if self.dop_prompt_embeds is not None:
self.dop_prompt_embeds = None self.dop_prompt_embeds = None
if self.dopsd_prompt_embeds is not None:
self.dopsd_prompt_embeds = None
def load_prompt_embedding(self, device=None): def load_prompt_embedding(self, device=None):
if not self.is_text_embedding_cached: if not self.is_text_embedding_cached:
@@ -2264,6 +2319,12 @@ class TextEmbeddingFileItemDTOMixin:
self.dop_prompt_embeds = self.prompt_embeds self.dop_prompt_embeds = self.prompt_embeds
else: else:
self.dop_prompt_embeds = PromptEmbeds.load(dop_path) self.dop_prompt_embeds = PromptEmbeds.load(dop_path)
if getattr(self, 'dopsd_self_ref', False) and self.dopsd_prompt_embeds is None:
if self._caption_was_dropped:
dopsd_path = self.get_dopsd_blank_text_embedding_path()
else:
dopsd_path = self.get_dopsd_text_embedding_path()
self.dopsd_prompt_embeds = PromptEmbeds.load(dopsd_path)
class TextEmbeddingCachingMixin: class TextEmbeddingCachingMixin:
def __init__(self: 'AiToolkitDataset', **kwargs): def __init__(self: 'AiToolkitDataset', **kwargs):
@@ -2402,6 +2463,50 @@ class TextEmbeddingCachingMixin:
prompt_embeds: PromptEmbeds = self.sd.encode_prompt(caption) prompt_embeds: PromptEmbeds = self.sd.encode_prompt(caption)
prompt_embeds.save(path) prompt_embeds.save(path)
del prompt_embeds del prompt_embeds
if getattr(file_item, 'dopsd_self_ref', False):
# D-OPSD teacher embeds encode with the item's own media as the reference
control_video_paths = getattr(file_item, 'control_video_paths', None) or []
if file_item.control_path is not None or len(control_video_paths) > 0:
raise ValueError(
"D-OPSD self-reference training cannot be combined with "
"control images/videos: the item itself must be the only "
f"reference. Offending item: {file_item.path}"
)
dopsd_targets = [(
file_item.get_dopsd_text_embedding_path(recalculate=True),
file_item.caption_dopsd,
)]
if self.dataset_config.caption_dropout_rate > 0:
dopsd_blank_path = file_item.get_dopsd_blank_text_embedding_path(recalculate=True)
if dopsd_blank_path != dopsd_targets[0][0]:
dopsd_targets.append((dopsd_blank_path, file_item.get_dopsd_dropout_caption()))
dopsd_targets = [t for t in dopsd_targets if not os.path.exists(t[0])]
if len(dopsd_targets) > 0:
if not did_move:
self.sd.set_device_state_preset('cache_text_encoder')
did_move = True
if file_item.is_video:
# own path rides through the video-ref presentation
ctrl_img = [file_item.path]
self.sd._ref_video_dataset_config = self.dataset_config
else:
# own bucketed pixels as the reference image
file_item.load_and_process_image(self.transform, only_load_latents=True)
img = file_item.tensor # (C, H, W) in [-1, 1]
ctrl_img = [
((img + 1.0) / 2.0)
.clamp(0, 1)
.unsqueeze(0)
.to(self.sd.device_torch, dtype=self.sd.torch_dtype)
]
# release the cached latent/pixels the load pulled in
file_item.cleanup()
if not self.sd.has_multiple_control_images:
ctrl_img = ctrl_img[0]
for path, caption in dopsd_targets:
prompt_embeds: PromptEmbeds = self.sd.encode_prompt(caption, control_images=ctrl_img)
prompt_embeds.save(path)
del prompt_embeds
file_item.is_text_embedding_cached = True file_item.is_text_embedding_cached = True
i += 1 i += 1
# restore device state # restore device state

View File

@@ -180,6 +180,10 @@ class BaseModel:
# control files may be VIDEOS (cached like dataset items, exposed on # control files may be VIDEOS (cached like dataset items, exposed on
# the batch as control_video_latents_list); see minimax_h3 ref2va # the batch as control_video_latents_list); see minimax_h3 ref2va
self.supports_video_control_images = False self.supports_video_control_images = False
# D-OPSD: cache per-item teacher text embeds (item's own media as reference 1)
self.dopsd_self_ref = False
# forces cache_tensors_to_disk on latent-caching datasets (BaseSDTrainProcess)
self.require_pixel_tensor_cache = False
# control images will come in as a list for encoding some things if true # control images will come in as a list for encoding some things if true
self.has_multiple_control_images = False self.has_multiple_control_images = False
# do not resize control images # do not resize control images

View File

@@ -218,6 +218,10 @@ class StableDiffusion:
# control files may be VIDEOS (paths exposed on the batch as # control files may be VIDEOS (paths exposed on the batch as
# control_video_paths_list); see minimax_h3 ref2va # control_video_paths_list); see minimax_h3 ref2va
self.supports_video_control_images = False self.supports_video_control_images = False
# D-OPSD: cache per-item teacher text embeds (item's own media as reference 1)
self.dopsd_self_ref = False
# forces cache_tensors_to_disk on latent-caching datasets (BaseSDTrainProcess)
self.require_pixel_tensor_cache = False
# control images will come in as a list for encoding some things if true # control images will come in as a list for encoding some things if true
self.has_multiple_control_images = False self.has_multiple_control_images = False
# do not resize control images # do not resize control images

View File

@@ -786,6 +786,7 @@ export const modelArchs: ModelArch[] = [
{ value: 'cg', label: 'Contrastive Guidance' }, { value: 'cg', label: 'Contrastive Guidance' },
{ value: 'ta', label: 'Training Adapter' }, { value: 'ta', label: 'Training Adapter' },
{ value: 'both', label: 'Contrastive Guidance + Training Adapter (default)' }, { value: 'both', label: 'Contrastive Guidance + Training Adapter (default)' },
{ value: 'none', label: 'None' },
], ],
getValue: (config: JobConfig) => { getValue: (config: JobConfig) => {
const assistantLoraPath = config?.config?.process?.[0]?.model?.assistant_lora_path; const assistantLoraPath = config?.config?.process?.[0]?.model?.assistant_lora_path;
@@ -797,7 +798,10 @@ export const modelArchs: ModelArch[] = [
if (hasAssistantLoraPath) { if (hasAssistantLoraPath) {
return 'ta'; return 'ta';
} }
return 'cg'; if (hasContrastiveGuidance) {
return 'cg';
}
return 'none';
}, },
onChange: (value: string, config: JobConfig, setJobConfig: (value: any, key: string) => void) => { onChange: (value: string, config: JobConfig, setJobConfig: (value: any, key: string) => void) => {
if (value === 'cg') { if (value === 'cg') {
@@ -822,6 +826,10 @@ export const modelArchs: ModelArch[] = [
if (!config?.config?.process?.[0]?.train?.guidance_loss_target) { if (!config?.config?.process?.[0]?.train?.guidance_loss_target) {
setJobConfig(3.5, 'config.process[0].train.guidance_loss_target'); setJobConfig(3.5, 'config.process[0].train.guidance_loss_target');
} }
} else if (value === 'none') {
setJobConfig(undefined, 'config.process[0].train.do_guidance_loss');
setJobConfig(undefined, 'config.process[0].train.guidance_loss_target');
setJobConfig(undefined, 'config.process[0].model.assistant_lora_path');
} }
}, },
doc: { doc: {
@@ -934,8 +942,13 @@ export const modelArchs: ModelArch[] = [
{ value: 'cg', label: 'Contrastive Guidance' }, { value: 'cg', label: 'Contrastive Guidance' },
{ value: 'ta', label: 'Training Adapter' }, { value: 'ta', label: 'Training Adapter' },
{ value: 'both', label: 'Contrastive Guidance + Training Adapter (default)' }, { value: 'both', label: 'Contrastive Guidance + Training Adapter (default)' },
{ value: 'dopsd', label: 'D-OPSD' },
{ value: 'none', label: 'None' },
], ],
getValue: (config: JobConfig) => { getValue: (config: JobConfig) => {
if (config?.config?.process?.[0]?.model?.model_kwargs?.dopsd) {
return 'dopsd';
}
const assistantLoraPath = config?.config?.process?.[0]?.model?.assistant_lora_path; const assistantLoraPath = config?.config?.process?.[0]?.model?.assistant_lora_path;
const hasAssistantLoraPath = assistantLoraPath && assistantLoraPath.trim() !== ''; const hasAssistantLoraPath = assistantLoraPath && assistantLoraPath.trim() !== '';
const hasContrastiveGuidance = config?.config?.process?.[0]?.train?.do_guidance_loss; const hasContrastiveGuidance = config?.config?.process?.[0]?.train?.do_guidance_loss;
@@ -945,9 +958,19 @@ export const modelArchs: ModelArch[] = [
if (hasAssistantLoraPath) { if (hasAssistantLoraPath) {
return 'ta'; return 'ta';
} }
return 'cg'; if (hasContrastiveGuidance) {
return 'cg';
}
return 'none';
}, },
onChange: (value: string, config: JobConfig, setJobConfig: (value: any, key: string) => void) => { onChange: (value: string, config: JobConfig, setJobConfig: (value: any, key: string) => void) => {
const kwargs = { ...(config?.config?.process?.[0]?.model?.model_kwargs ?? {}) };
if (value === 'dopsd') {
kwargs.dopsd = true;
} else {
delete kwargs.dopsd;
}
setJobConfig(kwargs, 'config.process[0].model.model_kwargs');
if (value === 'cg') { if (value === 'cg') {
setJobConfig(true, 'config.process[0].train.do_guidance_loss'); setJobConfig(true, 'config.process[0].train.do_guidance_loss');
setJobConfig(undefined, 'config.process[0].model.assistant_lora_path'); setJobConfig(undefined, 'config.process[0].model.assistant_lora_path');
@@ -970,6 +993,10 @@ export const modelArchs: ModelArch[] = [
if (!config?.config?.process?.[0]?.train?.guidance_loss_target) { if (!config?.config?.process?.[0]?.train?.guidance_loss_target) {
setJobConfig(3.5, 'config.process[0].train.guidance_loss_target'); setJobConfig(3.5, 'config.process[0].train.guidance_loss_target');
} }
} else if (value === 'dopsd' || value === 'none') {
setJobConfig(undefined, 'config.process[0].train.do_guidance_loss');
setJobConfig(undefined, 'config.process[0].train.guidance_loss_target');
setJobConfig(undefined, 'config.process[0].model.assistant_lora_path');
} }
}, },
doc: { doc: {
@@ -980,7 +1007,10 @@ export const modelArchs: ModelArch[] = [
break down. There are two different ways to train on this model without breaking the guidance break down. There are two different ways to train on this model without breaking the guidance
distillation: Contrastive Guidance and Training Adapter. Both have their pros and cons. The adapter is distillation: Contrastive Guidance and Training Adapter. Both have their pros and cons. The adapter is
faster, but will still break down over a long run. Contrastive Guidance is slower, but is less likely to faster, but will still break down over a long run. Contrastive Guidance is slower, but is less likely to
break down. break down. D-OPSD instead self-distills: a no-grad teacher pass sees the training target as its own
reference and its prediction becomes the training target for a reference-free pass, baking the reference
into your trigger word (or the caption itself when no trigger word is set). Re-caches latents with pixel
tensors.
</div> </div>
), ),
}, },
@@ -1009,12 +1039,12 @@ export const modelArchs: ModelArch[] = [
description: ( description: (
<div className="space-y-2"> <div className="space-y-2">
<p> <p>
How still-image references (dataset control images and sample ctrl images) are shown to the model. How still-image references (dataset control images and sample ctrl images) are shown to the model. Video
Video references always use the video path. references always use the video path.
</p> </p>
<p> <p>
<strong>Picture</strong>: the native ref2va recipe — a single-frame reference block, shown to Qwen3-VL as <strong>Picture</strong>: the native ref2va recipe — a single-frame reference block, shown to Qwen3-VL
a <code>&lt;Picture i&gt;</code> block, scaled down only. as a <code>&lt;Picture i&gt;</code> block, scaled down only.
</p> </p>
<p> <p>
<strong>Static video clip</strong>: the image is held for 5 frames (2 latent frames) and routed through <strong>Static video clip</strong>: the image is held for 5 frames (2 latent frames) and routed through
@@ -1035,10 +1065,10 @@ export const modelArchs: ModelArch[] = [
Reference-to-video: control images and videos condition the output as subject/style references (never as a Reference-to-video: control images and videos condition the output as subject/style references (never as a
first frame). References keep their own aspect and are matched to the target's pixel area (images scale down first frame). References keep their own aspect and are matched to the target's pixel area (images scale down
only, never up; a same-aspect video reference is exactly the target size). Each rides into the packed sequence only, never up; a same-aspect video reference is exactly the target size). Each rides into the packed sequence
as a reference block, and is also shown to the Qwen3-VL conditioner as a <code>&lt;Picture i&gt;</code> (image) as a reference block, and is also shown to the Qwen3-VL conditioner as a <code>&lt;Picture i&gt;</code>{' '}
or timestamped <code>&lt;Video k&gt;</code> (video) vision block. Training references come from the dataset (image) or timestamped <code>&lt;Video k&gt;</code> (video) vision block. Training references come from the
control path(s); sampling uses the sample ctrl images — always as references. The Image Reference Presentation dataset control path(s); sampling uses the sample ctrl images — always as references. The Image Reference
option can route still images through the video-reference path as short static clips. Presentation option can route still images through the video-reference path as short static clips.
</p> </p>
<p> <p>
Weights load like MiniMax-H3 (see that arch's notes) from the{' '} Weights load like MiniMax-H3 (see that arch's notes) from the{' '}

View File

@@ -1 +1 @@
VERSION = "0.12.26" VERSION = "0.12.27"