ImportError: cannot import name 'SpatialTransformer'

bazmatic commented 2 years ago

Throws this error when running the Training cell:

Traceback (most recent call last):
  File "/content/diffusers/examples/dreambooth/train_dreambooth.py", line 18, in <module>
    from diffusers import AutoencoderKL, DDPMScheduler, StableDiffusionPipeline, UNet2DConditionModel
  File "/usr/local/lib/python3.7/dist-packages/diffusers/__init__.py", line 21, in <module>
    from .models import AutoencoderKL, UNet2DConditionModel, UNet2DModel, VQModel
  File "/usr/local/lib/python3.7/dist-packages/diffusers/models/__init__.py", line 19, in <module>
    from .unet_2d import UNet2DModel
  File "/usr/local/lib/python3.7/dist-packages/diffusers/models/unet_2d.py", line 11, in <module>
    from .unet_blocks import UNetMidBlock2D, get_down_block, get_up_block
  File "/usr/local/lib/python3.7/dist-packages/diffusers/models/unet_blocks.py", line 20, in <module>
    from .attention import AttentionBlock, SpatialTransformer
ImportError: cannot import name 'SpatialTransformer' from 'diffusers.models.attention' (/usr/local/lib/python3.7/dist-packages/diffusers/models/attention.py)
Traceback (most recent call last):
  File "/usr/local/bin/accelerate", line 8, in <module>
    sys.exit(main())
  File "/usr/local/lib/python3.7/dist-packages/accelerate/commands/accelerate_cli.py", line 43, in main
    args.func(args)
  File "/usr/local/lib/python3.7/dist-packages/accelerate/commands/launch.py", line 837, in launch_command
    simple_launcher(args)
  File "/usr/local/lib/python3.7/dist-packages/accelerate/commands/launch.py", line 354, in simple_launcher
    raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
subprocess.CalledProcessError: Command '['/usr/bin/python3', '/content/diffusers/examples/dreambooth/train_dreambooth.py', '--image_captions_filename', '--save_starting_step=500', '--stop_text_encoder_training=1820', '--save_n_steps=500', '--Session_dir=/content/gdrive/MyDrive/Fast-Dreambooth/Sessions/bazmatix', '--pretrained_model_name_or_path=/content/stable-diffusion-v1-5', '--instance_data_dir=/content/gdrive/MyDrive/Fast-Dreambooth/Sessions/bazmatix/instance_images', '--output_dir=/content/models/bazmatix', '--instance_prompt=', '--seed=96576', '--resolution=512', '--mixed_precision=no', '--train_batch_size=1', '--gradient_accumulation_steps=1', '--use_8bit_adam', '--learning_rate=2e-6', '--lr_scheduler=polynomial', '--center_crop', '--lr_warmup_steps=0', '--max_train_steps=5200']' returned non-zero exit status 1.
Something went wrong

i2chris commented 2 years ago

I get exactly the same error.

This only happened after I chose the Premium GPU option in Colab and was given an A100:

!nvidia-smi -L GPU 0: A100-SXM4-40GB (UUID: GPU-fb627d06-76e2-7fe5-fa54-0811197d864f)

Training the unet...
Traceback (most recent call last):
  File "/content/diffusers/examples/dreambooth/train_dreambooth.py", line 18, in <module>
    from diffusers import AutoencoderKL, DDPMScheduler, StableDiffusionPipeline, UNet2DConditionModel
  File "/usr/local/lib/python3.7/dist-packages/diffusers/__init__.py", line 21, in <module>
    from .models import AutoencoderKL, UNet2DConditionModel, UNet2DModel, VQModel
  File "/usr/local/lib/python3.7/dist-packages/diffusers/models/__init__.py", line 19, in <module>
    from .unet_2d import UNet2DModel
  File "/usr/local/lib/python3.7/dist-packages/diffusers/models/unet_2d.py", line 11, in <module>
    from .unet_blocks import UNetMidBlock2D, get_down_block, get_up_block
  File "/usr/local/lib/python3.7/dist-packages/diffusers/models/unet_blocks.py", line 20, in <module>
    from .attention import AttentionBlock, SpatialTransformer
ImportError: cannot import name 'SpatialTransformer' from 'diffusers.models.attention' (/usr/local/lib/python3.7/dist-packages/diffusers/models/attention.py)
Traceback (most recent call last):
  File "/usr/local/bin/accelerate", line 8, in <module>
    sys.exit(main())
  File "/usr/local/lib/python3.7/dist-packages/accelerate/commands/accelerate_cli.py", line 43, in main
    args.func(args)
  File "/usr/local/lib/python3.7/dist-packages/accelerate/commands/launch.py", line 837, in launch_command
    simple_launcher(args)
  File "/usr/local/lib/python3.7/dist-packages/accelerate/commands/launch.py", line 354, in simple_launcher
    raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
subprocess.CalledProcessError: Command '['/usr/bin/python3', '/content/diffusers/examples/dreambooth/train_dreambooth.py', '--image_captions_filename', '--train_only_unet', '--Session_dir=/content/gdrive/MyDrive/Fast-Dreambooth/Sessions/theliperson', '--save_starting_step=1000', '--save_n_steps=0', '--pretrained_model_name_or_path=/content/stable-diffusion-v1-5', '--instance_data_dir=/content/gdrive/MyDrive/Fast-Dreambooth/Sessions/theliperson/instance_images', '--output_dir=/content/models/theliperson', '--instance_prompt=', '--seed=96576', '--resolution=512', '--mixed_precision=no', '--train_batch_size=1', '--gradient_accumulation_steps=1', '--use_8bit_adam', '--learning_rate=2e-6', '--lr_scheduler=polynomial', '--center_crop', '--lr_warmup_steps=0', '--max_train_steps=1500']' returned non-zero exit status 1.
Something went wrong

BempaOSD commented 2 years ago

I get the same error when using Colab Pro and Premium GPU-setting, but not when using 'regular' T4. I haven't had the time to dive deeper to the exact issue, but it seems related to GPU-type?

i2chris commented 2 years ago

I've gone back to the Standard GPU in Colab and it works fine.

GPU 0: Tesla T4 (UUID: GPU-7b3f2eed-e709-e150-0423-ca92a64f62a2)