monkey patch _compute_timestep_embedding

app.py

@@ -3,6 +3,7 @@ import numpy as np
 import random
 import spaces
 import torch
+import types
 from diffusers.pipelines.prx import PRXPipeline
 
 # monkey patch to add 1024 aspect ratios 
@@ -62,7 +63,7 @@ def get_timestep_embedding(
         emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
     return emb
  
-def _compute_timestep_embedding(timestep: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+def _compute_timestep_embedding(self, timestep: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
         return self.time_in(
             get_timestep_embedding(
                 timesteps=timestep,
@@ -111,7 +112,8 @@ pipe = PRXPipeline.from_pretrained(
     torch_dtype=dtype
 ).to(device)
 
-pipe.transformer._compute_timestep_embedding = _compute_timestep_embedding
+# Properly bind the method to the instance using types.MethodType
+pipe.transformer._compute_timestep_embedding = types.MethodType(_compute_timestep_embedding, pipe.transformer)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024