ModuleNotFoundError: No module named 'diffusers.models.unet_2d_condition'

[lousiming]

---

ModuleNotFoundError Traceback (most recent call last)
Cell In[2], line 3
1 import torch
----> 3 from distrifuser.pipelines import DistriSDXLPipeline
4 from distrifuser.utils import DistriConfig
6 distri_config = DistriConfig(height=1024, width=1024, warmup_steps=4)

File ~/miniconda3/envs/by/lib/python3.10/site-packages/distrifuser/pipelines.py:4
1 import torch
2 from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, UNet2DConditionModel
----> 4 from .models.distri_sdxl_unet_pp import DistriUNetPP
5 from .models.distri_sdxl_unet_tp import DistriUNetTP
6 from .models.naive_patch_sdxl import NaivePatchUNet

File ~/miniconda3/envs/by/lib/python3.10/site-packages/distrifuser/models/distri_sdxl_unet_pp.py:4
2 from diffusers import UNet2DConditionModel
3 from diffusers.models.attention_processor import Attention
----> 4 from diffusers.models.unet_2d_condition import UNet2DConditionOutput
5 from torch import distributed as dist, nn
7 from .base_model import BaseModel

ModuleNotFoundError: No module named 'diffusers.models.unet_2d_condition'

[lousiming]

pip install diffusers==0.31.0
I’m using this version

This is the one the code runs on

import torch

from distrifuser.pipelines import DistriSDXLPipeline
from distrifuser.utils import DistriConfig

distri_config = DistriConfig(height=1024, width=1024, warmup_steps=4)
pipeline = DistriSDXLPipeline.from_pretrained(
distri_config=distri_config,
pretrained_model_name_or_path="/home/eg4/lsm_test/models/stable-diffusion-xl-base-1.0",
variant="fp16",
use_safetensors=True,
)

pipeline.set_progress_bar_config(disable=distri_config.rank != 0)
image = pipeline(
prompt="Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
generator=torch.Generator(device="cuda").manual_seed(233),
).images[0]
if distri_config.rank == 0:
image.save("astronaut.png")