fix jit classmethod -> staticmethod

Author: Cui-yshoho

mindone/diffusers/models/attention_dispatch.py

@@ -158,37 +158,40 @@ class _AttentionBackendRegistry:
     _active_backend = AttentionBackendName(DIFFUSERS_ATTN_BACKEND)
     _checks_enabled = DIFFUSERS_ATTN_CHECKS
 
-    @classmethod
+    @staticmethod
     def register(
-        cls,
         backend: AttentionBackendName,
         constraints: Optional[List[Callable]] = None,
         supports_context_parallel: bool = False,
     ):
+        Registry = _AttentionBackendRegistry
         logger.debug(f"Registering attention backend: {backend} with constraints: {constraints}")
 
         def decorator(func):
-            cls._backends[backend] = func
-            cls._constraints[backend] = constraints or []
-            cls._supported_arg_names[backend] = set(inspect.signature(func).parameters.keys())
-            cls._supports_context_parallel[backend] = supports_context_parallel
+            Registry._backends[backend] = func
+            Registry._constraints[backend] = constraints or []
+            Registry._supported_arg_names[backend] = set(inspect.signature(func).parameters.keys())
+            Registry._supports_context_parallel[backend] = supports_context_parallel
             return func
 
         return decorator
 
-    @classmethod
-    def get_active_backend(cls):
-        return cls._active_backend, cls._backends[cls._active_backend]
+    @staticmethod
+    def get_active_backend():
+        Registry = _AttentionBackendRegistry
+        return Registry._active_backend, Registry._backends[Registry._active_backend]
 
-    @classmethod
-    def list_backends(cls):
-        return list(cls._backends.keys())
+    @staticmethod
+    def list_backends():
+        Registry = _AttentionBackendRegistry
+        return list(Registry._backends.keys())
 
-    @classmethod
+    @staticmethod
     def _is_context_parallel_enabled(
-        cls, backend: AttentionBackendName, parallel_config: Optional["ParallelConfig"]
+        backend: AttentionBackendName, parallel_config: Optional["ParallelConfig"]
     ) -> bool:
-        supports_context_parallel = backend in cls._supports_context_parallel
+        Registry = _AttentionBackendRegistry
+        supports_context_parallel = backend in Registry._supports_context_parallel
         is_degree_greater_than_1 = parallel_config is not None and (
             parallel_config.context_parallel_config.ring_degree > 1
             or parallel_config.context_parallel_config.ulysses_degree > 1

mindone/diffusers/models/transformers/transformer_skyreels_v2.py

@@ -114,8 +114,8 @@ def apply_rotary_emb(
                 out[..., 1::2] = x1 * sin + x2 * cos
                 return out.type_as(hidden_states)
 
-            query = apply_rotary_emb(query, rotary_emb)
-            key = apply_rotary_emb(key, rotary_emb)
+            query = apply_rotary_emb(query, *rotary_emb)
+            key = apply_rotary_emb(key, *rotary_emb)
 
         # I2V task
         hidden_states_img = None
@@ -420,13 +420,13 @@ def construct(self, hidden_states: ms.Tensor) -> ms.Tensor:
         freqs_cos = self.freqs_cos.split(split_sizes, dim=1)
         freqs_sin = self.freqs_sin.split(split_sizes, dim=1)
 
-        freqs_cos_f = freqs_cos[0][:ppf].view(ppf, 1, 1, -1).expand((ppf, pph, ppw, -1))
-        freqs_cos_h = freqs_cos[1][:pph].view(1, pph, 1, -1).expand((ppf, pph, ppw, -1))
-        freqs_cos_w = freqs_cos[2][:ppw].view(1, 1, ppw, -1).expand((ppf, pph, ppw, -1))
+        freqs_cos_f = freqs_cos[0][:ppf].view(ppf, 1, 1, -1).broadcast_to((ppf, pph, ppw, -1))
+        freqs_cos_h = freqs_cos[1][:pph].view(1, pph, 1, -1).broadcast_to((ppf, pph, ppw, -1))
+        freqs_cos_w = freqs_cos[2][:ppw].view(1, 1, ppw, -1).broadcast_to((ppf, pph, ppw, -1))
 
-        freqs_sin_f = freqs_sin[0][:ppf].view(ppf, 1, 1, -1).expand((ppf, pph, ppw, -1))
-        freqs_sin_h = freqs_sin[1][:pph].view(1, pph, 1, -1).expand((ppf, pph, ppw, -1))
-        freqs_sin_w = freqs_sin[2][:ppw].view(1, 1, ppw, -1).expand((ppf, pph, ppw, -1))
+        freqs_sin_f = freqs_sin[0][:ppf].view(ppf, 1, 1, -1).broadcast_to((ppf, pph, ppw, -1))
+        freqs_sin_h = freqs_sin[1][:pph].view(1, pph, 1, -1).broadcast_to((ppf, pph, ppw, -1))
+        freqs_sin_w = freqs_sin[2][:ppw].view(1, 1, ppw, -1).broadcast_to((ppf, pph, ppw, -1))
 
         freqs_cos = mint.cat([freqs_cos_f, freqs_cos_h, freqs_cos_w], dim=-1).reshape(1, ppf * pph * ppw, 1, -1)
         freqs_sin = mint.cat([freqs_sin_f, freqs_sin_h, freqs_sin_w], dim=-1).reshape(1, ppf * pph * ppw, 1, -1)