Update app.py

app.py

@@ -192,8 +192,15 @@ pipe.transformer.time_text_embed.fixed_text_proj = nn.Linear(3072, 256).to(devic
 def patched_time_embed(self, timestep, guidance, pooled_projections):
     # Compute the timestep embedding (expected shape: (B,256))
     time_out = self.time_proj(timestep)
-    # Use our fixed text projection (now with output dimension 256) on the pooled text embeddings.
-    text_out = self.fixed_text_proj(pooled_projections)
+    
+    # Ensure fixed_text_proj is set to map from 3072 to 256.
+    # If it doesn't exist or its output dimension is not 256, recreate it.
+    if (not hasattr(self, "fixed_text_proj")) or (self.fixed_text_proj.out_features != 256):
+        self.fixed_text_proj = nn.Linear(3072, 256).to(
+            device=pooled_projections.device, dtype=pooled_projections.dtype
+        )
+    
+    text_out = self.fixed_text_proj(pooled_projections)  # Should produce shape (B,256)
     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)