From 233e29225673ef8c978c35ab6aac130f7c6d6565 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Mon, 13 Apr 2026 09:37:34 -0600 Subject: [PATCH] Added some experimental low step things for zeta --- .../zeta_chroma/zeta_chroma_model.py | 3 + .../zeta_chroma/zeta_chroma_pipeline.py | 7 +- .../zeta_chroma/zeta_chroma_transformer.py | 4 + extensions_built_in/sd_trainer/SDTrainer.py | 8 +- jobs/process/BaseSDTrainProcess.py | 14 +- .../models/diffusion_feature_extraction.py | 166 +++++++++++++++++- 6 files changed, 194 insertions(+), 8 deletions(-) diff --git a/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py b/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py index bddbf09d..ca8b491f 100644 --- a/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py +++ b/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py @@ -241,6 +241,8 @@ class ZetaChromaModel(BaseModel): ): self.model.to(self.device_torch, dtype=self.torch_dtype) self.model.to(self.device_torch) + + do_low_step_schedule = gen_config.num_inference_steps <= 8 and gen_config.guidance_scale <= 1.0 sc = self.get_bucket_divisibility() gen_config.width = int(gen_config.width // sc * sc) @@ -256,6 +258,7 @@ class ZetaChromaModel(BaseModel): guidance_scale=gen_config.guidance_scale, latents=gen_config.latents, generator=generator, + low_step_schedule=do_low_step_schedule, **extra, ).images[0] return img diff --git a/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_pipeline.py b/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_pipeline.py index 2df2e6d6..8ea34931 100644 --- a/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_pipeline.py +++ b/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_pipeline.py @@ -10,6 +10,7 @@ from diffusers.utils import logging, replace_example_docstring from diffusers.pipelines.z_image.pipeline_output import ZImagePipelineOutput from extensions_built_in.diffusion_models.zeta_chroma.zeta_chroma_transformer import ( get_schedule, + get_low_step_schedule, prepare_latent_image_ids, make_text_position_ids, vae_unflatten, @@ -80,6 +81,7 @@ class ZetaChromaPipeline(ZImagePipeline): callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, callback_on_step_end_tensor_inputs: List[str] = ["latents"], max_sequence_length: int = 512, + low_step_schedule: bool = False, ): device = self._execution_device @@ -119,7 +121,10 @@ class ZetaChromaPipeline(ZImagePipeline): ) # --- Timestep schedule --- - timesteps = get_schedule(num_inference_steps, num_patches) + if low_step_schedule: + timesteps = get_low_step_schedule(num_inference_steps) + else: + timesteps = get_schedule(num_inference_steps, num_patches) # --- Denoising loop (CFG) --- img = noise diff --git a/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_transformer.py b/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_transformer.py index 450d8664..e2db51da 100644 --- a/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_transformer.py +++ b/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_transformer.py @@ -737,3 +737,7 @@ def get_schedule( mu = m * image_seq_len + b timesteps = time_shift(mu, 1.0, timesteps) return timesteps.tolist() + +def get_low_step_schedule(num_steps: int) -> list: + """Build uniform spaced timestep schedule from t=1 (noise) to t=0 (clean) to match training.""" + return torch.linspace(1, 0, num_steps + 1).tolist() \ No newline at end of file diff --git a/extensions_built_in/sd_trainer/SDTrainer.py b/extensions_built_in/sd_trainer/SDTrainer.py index 066a5eea..5f1bacda 100644 --- a/extensions_built_in/sd_trainer/SDTrainer.py +++ b/extensions_built_in/sd_trainer/SDTrainer.py @@ -355,7 +355,11 @@ class SDTrainer(BaseSDTrainProcess): vae = self.sd.vae # if not (self.model_config.arch in ["flux"]) or self.sd.vae.__class__.__name__ == "AutoencoderPixelMixer": # vae = self.sd.vae - self.dfe = load_dfe(self.train_config.diffusion_feature_extractor_path, vae=vae) + self.dfe = load_dfe( + self.train_config.diffusion_feature_extractor_path, + vae=vae, + sd=self.sd + ) self.dfe.to(self.device_torch) if hasattr(self.dfe, 'vision_encoder') and self.train_config.gradient_checkpointing: # must be set to train for gradient checkpointing to work @@ -660,7 +664,7 @@ class SDTrainer(BaseSDTrainProcess): dfe_loss += torch.nn.functional.mse_loss(pred_feature_list[i], target_feature_list[i], reduction="mean") additional_loss += dfe_loss * self.train_config.diffusion_feature_extractor_weight * 100.0 - elif self.dfe.version in [3, 4, 5, 6]: + elif self.dfe.version in [3, 4, 5, 6, 7]: dfe_loss = self.dfe( noise=noise, noise_pred=noise_pred, diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py index cd748c55..2ddf684f 100644 --- a/jobs/process/BaseSDTrainProcess.py +++ b/jobs/process/BaseSDTrainProcess.py @@ -1171,7 +1171,7 @@ class BaseSDTrainProcess(BaseTrainProcess): self.train_config.linear_timesteps, self.train_config.linear_timesteps2, self.train_config.timestep_type == 'linear', - self.train_config.timestep_type == 'one_step', + self.train_config.timestep_type in ['one_step', 'two_step', 'four_step', 'eight_step'], ]) timestep_type = 'linear' if linear_timesteps else None @@ -1225,8 +1225,16 @@ class BaseSDTrainProcess(BaseTrainProcess): if is_reg: content_or_style = self.train_config.content_or_style_reg - # if self.train_config.timestep_sampling == 'style' or self.train_config.timestep_sampling == 'content': - if self.train_config.timestep_type == 'next_sample': + if self.train_config.timestep_type in ['two_step', 'four_step', 'eight_step']: + if self.train_config.timestep_type == 'two_step': + indice_choices = [0, 499] + elif self.train_config.timestep_type == 'four_step': + indice_choices = [0, 250, 500, 750] + elif self.train_config.timestep_type == 'eight_step': + indice_choices = [0, 125, 250, 375, 500, 625, 750, 875] + timestep_indices = torch.tensor(random.choices(indice_choices, k=batch_size), device=self.device_torch) + timestep_indices = timestep_indices.long() + elif self.train_config.timestep_type == 'next_sample': timestep_indices = torch.randint( 0, num_train_timesteps - 2, # -1 for 0 idx, -1 so we can step diff --git a/toolkit/models/diffusion_feature_extraction.py b/toolkit/models/diffusion_feature_extraction.py index d50867c1..e9d8bdd2 100644 --- a/toolkit/models/diffusion_feature_extraction.py +++ b/toolkit/models/diffusion_feature_extraction.py @@ -7,9 +7,10 @@ import torch.nn.functional as F from diffusers import AutoencoderTiny from transformers import AutoImageProcessor, AutoModel, SiglipImageProcessor, SiglipVisionModel import lpips - +import weakref from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler +from toolkit.models.base_model import BaseModel class ResBlock(nn.Module): @@ -804,8 +805,165 @@ class DiffusionFeatureExtractor6(nn.Module): self.step += 1 return dino_loss + -def load_dfe(model_path, vae=None) -> DiffusionFeatureExtractor: +class DiffusionFeatureExtractor7(nn.Module): + def __init__(self, device=torch.device("cuda"), dtype=torch.bfloat16, vae=None, sd=None): + super().__init__() + self.version = 7 + self.sd_ref = weakref.ref(sd) if sd is not None else None + pretrained_model_name = "google/tipsv2-b14-dpt" + self.model = AutoModel.from_pretrained( + pretrained_model_name, + device_map=device, + dtype=torch.float32, + trust_remote_code=True + ).to(device) + + self.losses = {} + self.log_every = 100 + self.step = 0 + + def prepare_inputs(self, tensor_0_1: torch.Tensor): + """ + tensor_0_1: (bs, 3, h, w), float, values in [0, 1] + returns: {"pixel_values": (bs, 3, H, W)} ready for the vision transformer + """ + + if tensor_0_1.ndim != 4 or tensor_0_1.shape[1] != 3: + raise ValueError(f"Expected (bs, 3, h, w), got {tuple(tensor_0_1.shape)}") + + x = tensor_0_1.to(self.model.device, dtype=self.model.dtype) + + # Resize + # if not divisible by 16 or total pixels > max_res*max_res, resize to fit within 16 patches + max_res = 1024 + p = 14 + if (x.shape[-1] % p != 0) or (x.shape[-2] % p != 0) or (x.shape[-1] * x.shape[-2] > max_res * max_res): + target_h = x.shape[-2] + target_w = x.shape[-1] + if x.shape[-1] * target_h > max_res * max_res: + scale_factor = math.sqrt((max_res * max_res) / (target_w * target_h)) + target_h = int(target_h * scale_factor) + target_w = int(target_w * scale_factor) + target_h = (target_h // p) * p + target_w = (target_w // p) * p + x = F.interpolate(x, size=(target_h, target_w), mode="bilinear", align_corners=False) + + return x + + def forward( + self, + noise, + noise_pred, + noisy_latents, + timesteps, + batch: DataLoaderBatchDTO, + scheduler: CustomFlowMatchEulerDiscreteScheduler, + model=None + ): + dtype = torch.bfloat16 + device = self.sd_ref().vae.device + tensors = batch.tensor.to(device, dtype=dtype) + is_video = False + # stack time for video models on the batch dimension + if len(noise_pred.shape) == 5: + # B, C, T, H, W = images.shape + # only take first time + noise = noise[:, :, 0, :, :] + noise_pred = noise_pred[:, :, 0, :, :] + noisy_latents = noisy_latents[:, :, 0, :, :] + is_video = True + + if len(tensors.shape) == 5: + # batch is different + # (B, T, C, H, W) + # only take first time + tensors = tensors[:, 0, :, :, :] + + with torch.no_grad(): + tv = timesteps.to(noise_pred.device).to(noise_pred.dtype) / 1000.0 + # expand shape to match noise_pred + while len(tv.shape) < len(noise_pred.shape): + tv = tv.unsqueeze(-1) + # min 0.001 + tv = torch.clamp(tv, min=0.001) + + # step latent + x0 = noisy_latents - tv * noise_pred + + stepped_latents = x0 + + latents = stepped_latents.to(self.sd_ref().vae.device, dtype=self.sd_ref().vae.dtype) + + tensors_n1p1 = self.sd_ref().decode_latents(latents) + + pred_images = (tensors_n1p1 + 1) / 2 # 0 to 1 + + device = self.model.device + dtype = self.model.dtype + + with torch.no_grad(): + target_img = tensors.to(device, dtype=dtype) + # go from -1 to 1 to 0 to 1 + target_img = (target_img + 1) / 2 + target = self.prepare_inputs(target_img) + target = self.model(target) + + pred_images = pred_images.to(device, dtype=dtype) + pred = self.prepare_inputs(pred_images) + pred = self.model(pred) + + + depth_loss = torch.nn.functional.l1_loss( + pred.depth.float(), target.depth.float() + ) + + normals_loss = torch.nn.functional.l1_loss( + pred.normals.float(), target.normals.float() + ) + + segmentation_loss = torch.nn.functional.l1_loss( + pred.segmentation.float(), target.segmentation.float() + ) + + total_loss = (depth_loss + normals_loss + segmentation_loss) / 3.0 + + if 'total' not in self.losses: + self.losses['total'] = total_loss.item() + else: + self.losses['total'] += total_loss.item() + + if 'depth' not in self.losses: + self.losses['depth'] = depth_loss.item() + else: + self.losses['depth'] += depth_loss.item() + + if 'normals' not in self.losses: + self.losses['normals'] = normals_loss.item() + else: + self.losses['normals'] += normals_loss.item() + + if 'segmentation' not in self.losses: + self.losses['segmentation'] = segmentation_loss.item() + else: + self.losses['segmentation'] += segmentation_loss.item() + + with torch.no_grad(): + if self.step % self.log_every == 0 and self.step > 0: + print(f"DFE losses:") + for key in self.losses: + self.losses[key] /= self.log_every + # print in 2.000e-01 format + print(f" - {key}: {self.losses[key]:.3e}") + self.losses[key] = 0.0 + + # total_loss += mse_loss + self.step += 1 + + return total_loss + +def load_dfe(model_path, vae=None, sd: 'BaseModel' = None) -> DiffusionFeatureExtractor: if model_path == "v3": dfe = DiffusionFeatureExtractor3(vae=vae) dfe.eval() @@ -822,6 +980,10 @@ def load_dfe(model_path, vae=None) -> DiffusionFeatureExtractor: dfe = DiffusionFeatureExtractor6(vae=vae) dfe.eval() return dfe + if model_path == "v7": + dfe = DiffusionFeatureExtractor7(vae=vae, sd=sd) + dfe.eval() + return dfe if not os.path.exists(model_path): raise FileNotFoundError(f"Model file not found: {model_path}") # if it ende with safetensors