upgrade diffusers version to 0.31.0

.gitignore

@@ -2,4 +2,5 @@
 .DS_Store
 build
 dist
-*.egg-info
+*.egg-info
+*__pycache__*

distrifuser/models/distri_sdxl_unet_pp.py

@@ -1,7 +1,7 @@
 import torch
 from diffusers import UNet2DConditionModel
 from diffusers.models.attention_processor import Attention
-from diffusers.models.unet_2d_condition import UNet2DConditionOutput
+from diffusers.models.unets.unet_2d_condition import UNet2DConditionOutput
 from torch import distributed as dist, nn
 
 from .base_model import BaseModel

distrifuser/models/distri_sdxl_unet_tp.py

@@ -2,7 +2,7 @@
 from diffusers import UNet2DConditionModel
 from diffusers.models.attention import Attention, FeedForward
 from diffusers.models.resnet import ResnetBlock2D
-from diffusers.models.unet_2d_condition import UNet2DConditionOutput
+from diffusers.models.unets.unet_2d_condition import UNet2DConditionOutput
 from torch import distributed as dist, nn
 
 from distrifuser.modules.base_module import BaseModule
@@ -157,7 +157,11 @@ def forward(
                 if self.buffer_list is None:
                     self.buffer_list = [torch.empty_like(output) for _ in range(2)]
                 dist.all_gather(
-                    self.buffer_list, output.contiguous(), group=distri_config.split_group(), async_op=False
+                    self.buffer_list, output.contiguous(), 
+                    group=distri_config.split_group, 
+                    # original code
+                    # group=distri_config.split_group(), 
+                    async_op=False
                 )
                 torch.cat(self.buffer_list, dim=0, out=self.output_buffer)
                 output = self.output_buffer

distrifuser/models/naive_patch_sdxl.py

@@ -1,6 +1,6 @@
 import torch
 from diffusers import UNet2DConditionModel
-from diffusers.models.unet_2d_condition import UNet2DConditionOutput
+from diffusers.models.unets.unet_2d_condition import UNet2DConditionOutput
 from torch import distributed as dist
 
 from .base_model import BaseModel

distrifuser/modules/pp/attn.py

@@ -64,9 +64,13 @@ def forward(
             hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
         )
 
-        args = () if USE_PEFT_BACKEND else (scale,)
-        query = attn.to_q(hidden_states, *args)
-
+        # Handle scale parameter based on PEFT backend
+        if USE_PEFT_BACKEND:
+            query = attn.to_q(hidden_states)
+        else:
+            query = attn.to_q(hidden_states)
+            query = query * scale
+
         if encoder_hidden_states is None:
             encoder_hidden_states = hidden_states
 
@@ -117,8 +121,14 @@ def _forward(self, hidden_states: torch.FloatTensor, scale: float = 1.0):
 
         batch_size, sequence_length, _ = hidden_states.shape
 
-        args = () if USE_PEFT_BACKEND else (scale,)
-        query = attn.to_q(hidden_states, *args)
+        # args = () if USE_PEFT_BACKEND else (scale,)
+        # query = attn.to_q(hidden_states, *args)
+
+        if USE_PEFT_BACKEND:
+            query = attn.to_q(hidden_states)
+        else:
+            query = attn.to_q(hidden_states)
+            query = query * scale
 
         encoder_hidden_states = hidden_states
 
@@ -156,7 +166,13 @@ def _forward(self, hidden_states: torch.FloatTensor, scale: float = 1.0):
         hidden_states = hidden_states.to(query.dtype)
 
         # linear proj
-        hidden_states = attn.to_out[0](hidden_states, *args)
+        # hidden_states = attn.to_out[0](hidden_states, *args)
+
+        if USE_PEFT_BACKEND:
+            hidden_states = attn.to_out[0](hidden_states)
+        else:
+            hidden_states = attn.to_out[0](hidden_states)
+            hidden_states = hidden_states * scale
         # dropout
         hidden_states = attn.to_out[1](hidden_states)
 

distrifuser/modules/tp/resnet.py

@@ -1,5 +1,6 @@
 import torch.cuda
-from diffusers.models.resnet import Downsample2D, ResnetBlock2D, Upsample2D, USE_PEFT_BACKEND
+from diffusers.models.resnet import Downsample2D, ResnetBlock2D, Upsample2D
+from diffusers.utils import USE_PEFT_BACKEND
 from torch import distributed as dist
 from torch import nn
 from torch.nn import functional as F
@@ -192,7 +193,8 @@ def forward(
 
         if module.conv_shortcut is not None:
             input_tensor = (
-                module.conv_shortcut(input_tensor, scale) if not USE_PEFT_BACKEND else self.conv_shortcut(input_tensor)
+                # module.conv_shortcut(input_tensor, scale)
+                module.conv_shortcut(input_tensor) * scale if not USE_PEFT_BACKEND else self.conv_shortcut(input_tensor)
             )
 
         output_tensor = (input_tensor + hidden_states) / module.output_scale_factor

distrifuser/pipelines.py

@@ -37,7 +37,10 @@ def from_pretrained(distri_config: DistriConfig, **kwargs):
             raise ValueError(f"Unknown parallelism: {distri_config.parallelism}")
 
         pipeline = StableDiffusionXLPipeline.from_pretrained(
-            pretrained_model_name_or_path, torch_dtype=torch_dtype, unet=unet, **kwargs
+            pretrained_model_name_or_path, 
+            torch_dtype=torch_dtype, 
+            unet=unet, 
+            **kwargs
         ).to(device)
         return DistriSDXLPipeline(pipeline, distri_config)
 

setup.py

@@ -12,7 +12,7 @@
         author="Muyang Li, Tianle Cai, Jiaxin Cao, Qinsheng Zhang, Han Cai, Junjie Bai, Yangqing Jia, Ming-Yu Liu, Kai Li and Song Han",
         author_email="muyangli@mit.edu",
         packages=find_packages(),
-        install_requires=["torch>=2.2", "diffusers==0.24.0", "transformers", "tqdm"],
+        install_requires=["torch>=2.2", "diffusers>=0.31.0", "transformers", "tqdm"],
         url="https://github.com/mit-han-lab/distrifuser",
         description="DistriFusion: Distributed Parallel Inference for High-Resolution Diffusion Models",
         long_description=long_description,