Minor fix in AuraFlow callback

Skquark · Aug 3, 2024 · 56a21c5 · 56a21c5
1 parent bfe3f6a
commit 56a21c5
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py b/src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py
@@ -124,6 +124,9 @@ class AuraFlowPipeline(DiffusionPipeline):
     """
 
     _optional_components = []
+    _callback_tensor_inputs = [
+        "latents",
+    ]
     model_cpu_offload_seq = "text_encoder->transformer->vae"
 
     def __init__(