From 92f831eb242ad059bf1446daefd75090e14f2327 Mon Sep 17 00:00:00 2001
From: Alexander Piskun <bigcat88@icloud.com>
Date: Sat, 11 Jan 2025 22:14:26 +0300
Subject: [PATCH] rebased onto "master"

---
 comfy/ldm/cosmos/model.py | 26 +++++++++++++-------------
 pyproject.toml            |  1 +
 2 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/comfy/ldm/cosmos/model.py b/comfy/ldm/cosmos/model.py
index 05dd3846..b8135ea9 100644
--- a/comfy/ldm/cosmos/model.py
+++ b/comfy/ldm/cosmos/model.py
@@ -217,19 +217,19 @@ class GeneralDIT(nn.Module):
             raise ValueError(f"Unknown pos_emb_cls {self.pos_emb_cls}")
 
         logging.debug(f"Building positional embedding with {self.pos_emb_cls} class, impl {cls_type}")
-        kwargs = dict(
-            model_channels=self.model_channels,
-            len_h=self.max_img_h // self.patch_spatial,
-            len_w=self.max_img_w // self.patch_spatial,
-            len_t=self.max_frames // self.patch_temporal,
-            is_learnable=self.pos_emb_learnable,
-            interpolation=self.pos_emb_interpolation,
-            head_dim=self.model_channels // self.num_heads,
-            h_extrapolation_ratio=self.rope_h_extrapolation_ratio,
-            w_extrapolation_ratio=self.rope_w_extrapolation_ratio,
-            t_extrapolation_ratio=self.rope_t_extrapolation_ratio,
-            device=device,
-        )
+        kwargs = {
+            "model_channels": self.model_channels,
+            "len_h": self.max_img_h // self.patch_spatial,
+            "len_w": self.max_img_w // self.patch_spatial,
+            "len_t": self.max_frames // self.patch_temporal,
+            "is_learnable": self.pos_emb_learnable,
+            "interpolation": self.pos_emb_interpolation,
+            "head_dim": self.model_channels // self.num_heads,
+            "h_extrapolation_ratio": self.rope_h_extrapolation_ratio,
+            "w_extrapolation_ratio": self.rope_w_extrapolation_ratio,
+            "t_extrapolation_ratio": self.rope_t_extrapolation_ratio,
+            "device": device,
+        }
         self.pos_embedder = cls_type(
             **kwargs,
         )
diff --git a/pyproject.toml b/pyproject.toml
index b747d6ef..296bc45d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -19,5 +19,6 @@ lint.select = [
   # The "F" series in Ruff stands for "Pyflakes" rules, which catch various Python syntax errors and undefined names.
   # See all rules here: https://docs.astral.sh/ruff/rules/#pyflakes-f
   "F",
+  "C408",  # unnecessary dict(), list() or tuple() calls that can be rewritten as empty literals.
 ]
 exclude = ["*.ipynb"]