Update app.py

app.py

@@ -186,15 +186,15 @@ pipe.encode_prompt = custom_encode_prompt.__get__(pipe)
 
 pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
 
-def patched_time_embed(x):
-    # Get the weight and bias from the time embedder.
-    weight = pipe.transformer.time_text_embed.weight  # shape: (768, 3072)
-    bias = pipe.transformer.time_text_embed.bias
-    # Transpose weight to shape (3072, 768) so that input (B, 3072) can multiply correctly.
-    return F.linear(x, weight.t(), bias)
-
-# Apply the patch:
-pipe.transformer.time_text_embed.forward = patched_time_embed
+def patched_time_embed(self, timestep, guidance, pooled_projections):
+    # 'timestep' and 'guidance' are not used in our patch; we operate on pooled_projections.
+    weight = self.weight  # shape should be (768, 3072)
+    bias = self.bias      # shape (768,)
+    # Transpose weight so that pooled_projections (shape (B, 3072)) multiplies correctly:
+    return F.linear(pooled_projections, weight.t(), bias)
+    
+# Patch the time embedder on the transformer:
+pipe.transformer.time_text_embed.forward = patched_time_embed.__get__(pipe.transformer.time_text_embed)
 # History functions
 def append_to_history(image, prompt, seed, width, height, guidance_scale, steps, history):
     """Store only the final generated image"""