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