Update app.py

app.py

@@ -187,24 +187,14 @@ 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)
 
+pipe.transformer.time_text_embed.fixed_text_proj = nn.Linear(3072, 256).to(device, dtype=dtype)
 
 def patched_time_embed(self, timestep, guidance, pooled_projections):
     # Compute the timestep embedding (expected shape: (B,256))
     time_out = self.time_proj(timestep)
-    
-    # On first call, attach a new linear layer to project pooled_projections from 3072 to 256,
-    # matching the dimension of time_out.
-    if not hasattr(self, "fixed_text_proj"):
-        self.fixed_text_proj = nn.Linear(3072, 256).to(
-            device=pooled_projections.device, dtype=pooled_projections.dtype
-        )
-    
-    # Apply the new projection to pooled text embeddings.
-    text_out = self.fixed_text_proj(pooled_projections)  # now shape (B,256)
-    
+    # Use our fixed text projection (now with output dimension 256) on the pooled text embeddings.
+    text_out = self.fixed_text_proj(pooled_projections)
     return time_out + text_out
-
-
 # Apply the patch after the pipeline is created and patched with your custom encode methods:
 pipe.transformer.time_text_embed.forward = patched_time_embed.__get__(pipe.transformer.time_text_embed)