add cache

src/pipeline.py

@@ -656,7 +656,30 @@ 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():
+            attn_processor.bank_kv.clear()
+        # cache with warmup latents
+        start_idx = latents.shape[1] - 32
+        warmup_latents = latents[:, start_idx:, :]
+        warmup_latent_ids = latent_image_ids[start_idx:, :]
+        t = torch.tensor([timesteps[0]], device=device)
+        timestep = t.expand(warmup_latents.shape[0]).to(latents.dtype)
+        _ = self.transformer(
+                    hidden_states=warmup_latents,  
+                    cond_hidden_states=cond_latents,
+                    timestep=timestep/ 1000,
+                    guidance=guidance,
+                    pooled_projections=pooled_prompt_embeds,
+                    encoder_hidden_states=prompt_embeds,
+                    txt_ids=text_ids,
+                    img_ids=warmup_latent_ids,
+                    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):