support cfg-zero*

src/pipeline.py

@@ -526,9 +526,11 @@ class FluxPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FromSingleFileMixin):
             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,
-            spatial_images=None,
-            subject_images=None,
+            spatial_images=[],
+            subject_images=[],
             cond_size=512,
+            use_zero_init: Optional[bool] = True,
+            zero_steps: Optional[int] = 0,
     ):
 
         height = height or self.default_sample_size * self.vae_scale_factor
@@ -656,7 +658,7 @@ class FluxPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FromSingleFileMixin):
             guidance = guidance.expand(latents.shape[0])
         else:
             guidance = None
-        
+
         ## Caching conditions
         # clean the cache
         for name, attn_processor in self.transformer.attn_processors.items():
@@ -679,7 +681,7 @@ class FluxPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FromSingleFileMixin):
                     joint_attention_kwargs=self.joint_attention_kwargs,
                     return_dict=False,
                 )[0]
-        
+
         # 6. Denoising loop
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
@@ -700,6 +702,9 @@ class FluxPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FromSingleFileMixin):
                     joint_attention_kwargs=self.joint_attention_kwargs,
                     return_dict=False,
                 )[0]
+                
+                if (i <= zero_steps) and use_zero_init:
+                    noise_pred = noise_pred*0.
 
                 # compute the previous noisy sample x_t -> x_t-1
                 latents_dtype = latents.dtype
@@ -742,4 +747,4 @@ class FluxPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FromSingleFileMixin):
         if not return_dict:
             return (image,)
 
-        return FluxPipelineOutput(images=image)
+        return FluxPipelineOutput(images=image)