Added some experimental low step things for zeta

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

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

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

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,