diffsynth-engine 0.7.1.dev3__py3-none-any.whl → 0.7.1.dev5__py3-none-any.whl

diffsynth_engine/models/qwen_image/qwen_image_dit_nunchaku.py

@@ -179,6 +179,7 @@ class QwenImageTransformerBlockNunchaku(QwenImageTransformerBlock):
         rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
         attn_mask: Optional[torch.Tensor] = None,
         attn_kwargs: Optional[Dict[str, Any]] = None,
+        modulate_index: Optional[List[int]] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         if self.use_nunchaku_awq:
             img_mod_params = self.img_mod(temb)  # [B, 6*dim]

diffsynth_engine/tools/qwen_image_upscaler_tool.py

@@ -12,7 +12,7 @@ from diffsynth_engine.configs import QwenImagePipelineConfig
 from diffsynth_engine.pipelines.qwen_image import QwenImagePipeline
 from diffsynth_engine.models.qwen_image import QwenImageVAE
 from diffsynth_engine.models.basic.lora import LoRALinear
-from diffsynth_engine.models.qwen_image.qwen_image_dit import QwenImageTransformerBlock
+from diffsynth_engine.models.qwen_image.qwen_image_dit import QwenImageTransformerBlock, QwenEmbedRope
 from diffsynth_engine.utils import logging
 from diffsynth_engine.utils.loader import load_file
 from diffsynth_engine.utils.download import fetch_model
@@ -32,6 +32,7 @@ def odtsr_forward():
     """
     original_lora_forward = LoRALinear.forward
     original_modulate = QwenImageTransformerBlock._modulate
+    original_rope_forward = QwenEmbedRope.forward
 
     def lora_batch_cfg_forward(self, x):
         y = nn.Linear.forward(self, x)
@@ -50,6 +51,49 @@ def odtsr_forward():
                 y[:, L:] += lora(x2)
             return y
 
+    def optimized_rope_forward(self, video_fhw, txt_length, device):
+        if self.pos_freqs.device != device:
+            self.pos_freqs = self.pos_freqs.to(device)
+            self.neg_freqs = self.neg_freqs.to(device)
+
+        vid_freqs = []
+        max_vid_index = 0
+        idx = 0
+        for fhw in video_fhw:
+            frame, height, width = fhw
+            rope_key = f"{idx}_{height}_{width}"
+
+            if rope_key not in self.rope_cache:
+                seq_lens = frame * height * width
+                freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1)
+                freqs_neg = self.neg_freqs.split([x // 2 for x in self.axes_dim], dim=1)
+                freqs_frame = freqs_pos[0][idx : idx + frame].view(frame, 1, 1, -1).expand(frame, height, width, -1)
+                if self.scale_rope:
+                    freqs_height = torch.cat(
+                        [freqs_neg[1][-(height - height // 2) :], freqs_pos[1][: height // 2]], dim=0
+                    )
+                    freqs_height = freqs_height.view(1, height, 1, -1).expand(frame, height, width, -1)
+                    freqs_width = torch.cat([freqs_neg[2][-(width - width // 2) :], freqs_pos[2][: width // 2]], dim=0)
+                    freqs_width = freqs_width.view(1, 1, width, -1).expand(frame, height, width, -1)
+
+                else:
+                    freqs_height = freqs_pos[1][:height].view(1, height, 1, -1).expand(frame, height, width, -1)
+                    freqs_width = freqs_pos[2][:width].view(1, 1, width, -1).expand(frame, height, width, -1)
+
+                freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape(seq_lens, -1)
+                self.rope_cache[rope_key] = freqs.clone().contiguous()
+            vid_freqs.append(self.rope_cache[rope_key])
+            if self.scale_rope:
+                max_vid_index = max(height // 2, width // 2, max_vid_index)
+            else:
+                max_vid_index = max(height, width, max_vid_index)
+
+        txt_freqs = self.pos_freqs[max_vid_index : max_vid_index + txt_length, ...]
+        vid_freqs = torch.cat(vid_freqs, dim=0)
+
+        return vid_freqs, txt_freqs
+
+
     def optimized_modulate(self, x, mod_params, index=None):
         if mod_params.ndim == 2:
             shift, scale, gate = mod_params.chunk(3, dim=-1)
@@ -72,12 +116,14 @@ def odtsr_forward():
 
     LoRALinear.forward = lora_batch_cfg_forward
     QwenImageTransformerBlock._modulate = optimized_modulate
+    QwenEmbedRope.forward = optimized_rope_forward
 
     try:
         yield
     finally:
         LoRALinear.forward = original_lora_forward
         QwenImageTransformerBlock._modulate = original_modulate
+        QwenEmbedRope.forward = original_rope_forward
 
 
 class QwenImageUpscalerTool:

{diffsynth_engine-0.7.1.dev3.dist-info → diffsynth_engine-0.7.1.dev5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsynth_engine
-Version: 0.7.1.dev3
+Version: 0.7.1.dev5
 Author: MuseAI x ModelScope
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent

{diffsynth_engine-0.7.1.dev3.dist-info → diffsynth_engine-0.7.1.dev5.dist-info}/RECORD

@@ -122,7 +122,7 @@ diffsynth_engine/models/qwen_image/__init__.py,sha256=_6f0LWaoLdDvD2CsjK2OzEIQry
 diffsynth_engine/models/qwen_image/qwen2_5_vl.py,sha256=Eu-r-c42t_q74Qpwz21ToCGHpvSi7VND4B1EI0e-ePA,57748
 diffsynth_engine/models/qwen_image/qwen_image_dit.py,sha256=mMU4zeZi8-uJe9voznNIxZCTCqJPbPXkMxHwgcqJ6z8,24640
 diffsynth_engine/models/qwen_image/qwen_image_dit_fbcache.py,sha256=LIv9X_BohKk5rcEzyl3ATLwd8MSoFX43wjkArQ68nq8,4828
-diffsynth_engine/models/qwen_image/qwen_image_dit_nunchaku.py,sha256=1y1BkPRrX4_RioKjM09D9f9PK9neug1nSGJka0D9bvM,13516
+diffsynth_engine/models/qwen_image/qwen_image_dit_nunchaku.py,sha256=EIojuf27haxqI4wkJE_Y17HMjP82-iqvyJ5v5Kjns3o,13568
 diffsynth_engine/models/qwen_image/qwen_image_vae.py,sha256=FpauZV9IVvpvBeS9volu7kzH2mmCISS86AbHt0Jk2bQ,38442
 diffsynth_engine/models/sd/__init__.py,sha256=hjoKRnwoXOLD0wude-w7I6wK5ak7ACMbnbkPuBB2oU0,380
 diffsynth_engine/models/sd/sd_controlnet.py,sha256=kMGfIdriXhC7reT6iO2Z0rPICXEkXpytjeBQcR_sjT8,50577
@@ -186,7 +186,7 @@ diffsynth_engine/tools/flux_inpainting_tool.py,sha256=qHsYKUG20A19ujRdocpIPC4a_H
 diffsynth_engine/tools/flux_outpainting_tool.py,sha256=ff4qUj2mMYW6GMts7ifnJG7Rth55pfuggopRCyAXwJ8,3894
 diffsynth_engine/tools/flux_reference_tool.py,sha256=6v0NRZPsDEHFlPruO-ZJTB4rYWxKVAlmnYEeandD3r8,4723
 diffsynth_engine/tools/flux_replace_tool.py,sha256=AOyEGxHsaNwpTS2VChAieIfECgMxlKsRw0lWPm1k9C0,4627
-diffsynth_engine/tools/qwen_image_upscaler_tool.py,sha256=TFtITz113zoqsdRibVuLtWF8JEhGTqzyV2ZGHJuuYKw,13876
+diffsynth_engine/tools/qwen_image_upscaler_tool.py,sha256=GMhV7Sphg2zgkOJhnZeLVWQJQv1d6QnOuQZXEvHgIyI,16222
 diffsynth_engine/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diffsynth_engine/utils/cache.py,sha256=Ivef22pCuhEq-4H00gSvkLS8ceVZoGis7OSitYL6gH4,2101
 diffsynth_engine/utils/constants.py,sha256=Tsn3EAByfZra-nGcx0NEcP9nWTPKaDGdatosE3BuPGE,3846
@@ -209,8 +209,8 @@ diffsynth_engine/utils/video.py,sha256=8FCaeqIdUsWMgWI_6SO9SPynsToGcLCQAVYFTc4CD
 diffsynth_engine/utils/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diffsynth_engine/utils/memory/linear_regression.py,sha256=oW_EQEw13oPoyUrxiL8A7Ksa5AuJ2ynI2qhCbfAuZbg,3930
 diffsynth_engine/utils/memory/memory_predcit_model.py,sha256=EXprSl_zlVjgfMWNXP-iw83Ot3hyMcgYaRPv-dvyL84,3943
-diffsynth_engine-0.7.1.dev3.dist-info/licenses/LICENSE,sha256=x7aBqQuVI0IYnftgoTPI_A0I_rjdjPPQkjnU6N2nikM,11346
-diffsynth_engine-0.7.1.dev3.dist-info/METADATA,sha256=GdfffMwz8CD9vSlEKGlzjwp_fO19sYw0ulei0vx6rQY,1163
-diffsynth_engine-0.7.1.dev3.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-diffsynth_engine-0.7.1.dev3.dist-info/top_level.txt,sha256=6zgbiIzEHLbhgDKRyX0uBJOV3F6VnGGBRIQvSiYYn6w,17
-diffsynth_engine-0.7.1.dev3.dist-info/RECORD,,
+diffsynth_engine-0.7.1.dev5.dist-info/licenses/LICENSE,sha256=x7aBqQuVI0IYnftgoTPI_A0I_rjdjPPQkjnU6N2nikM,11346
+diffsynth_engine-0.7.1.dev5.dist-info/METADATA,sha256=76gzYfIIeo_71jVybkzGLWiMpkm95ifPNZkL12gCRj8,1163
+diffsynth_engine-0.7.1.dev5.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+diffsynth_engine-0.7.1.dev5.dist-info/top_level.txt,sha256=6zgbiIzEHLbhgDKRyX0uBJOV3F6VnGGBRIQvSiYYn6w,17
+diffsynth_engine-0.7.1.dev5.dist-info/RECORD,,