KeyError in visualize attn maps

[OchiaiNao]

I am using the notebook given #54
When I run the loading part:

ldm_stable.load_model('/content/drive/MyDrive/diffusion/logs/<painting>.bin')

Returns this:

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
[<ipython-input-12-73c40f26b44b>](https://localhost:8080/#) in <cell line: 11>()
      9 model_id = "CompVis/stable-diffusion-v1-4"
     10 ldm_stable = CustomDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", scheduler=scheduler).to(device)
---> 11 ldm_stable.load_model('/content/drive/MyDrive/diffusion/logs/<painting>.bin')
     12 tokenizer = ldm_stable.tokenizer

[/content/prompt-to-prompt/custom-diffusion/src/diffusers_model_pipeline.py](https://localhost:8080/#) in load_model(self, save_path, compress)
    495                 if compress and ('to_k' in name or 'to_v' in name):
    496                     params.data += st['unet'][name]['u']@st['unet'][name]['v']
--> 497                 elif name in st['unet']:
    498                     params.data.copy_(st['unet'][f'{name}'])

KeyError: 'unet'

How can I solve this?

[nupurkmr9]

Hi, can you mention the exact training command used to train the model .bin? Thanks.

[OchiaiNao]

Hi, this is my test train code:

accelerate launch train_custom_diffusion.py \
  --pretrained_model_name_or_path=CompVis/stable-diffusion-v1-4 \
  --instance_data_dir=/content/drive/MyDrive/diffusion/class_woman/young_woman \
  --output_dir=/content/drive/MyDrive/diffusion/train_logs/young_woman \
  --class_data_dir=/content/drive/MyDrive/diffusion/real_woman/ \
  --with_prior_preservation --real_prior --prior_loss_weight=1.0 \
  --instance_prompt="photo of a <young> woman"  \
  --class_prompt='woman' --num_class_images 150 \
  --resolution=512  \
  --train_batch_size=2  \
  --learning_rate=5e-6  \
  --lr_warmup_steps=500 \
  --max_train_steps=2000 \
  --scale_lr --hflip \
  --modifier_token "<young>" \
  --validation_prompt="photo of a <young> woman" \
  --no_safe_serialization \
  --enable_xformers_memory_efficient_attention \
  --report_to="wandb"`

[nupurkmr9]

Hi if you are using this to train the model, you would need to update the model load command as below:

ldm_stable = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16).to("cuda")
ldm_stable.unet.load_attn_procs("path-to-save-model", weight_name="pytorch_custom_diffusion_weights.bin")
ldm_stable.load_textual_inversion("path-to-save-model", weight_name="<young>.bin")

[OchiaiNao]

Hi, I used what you mentioned to load new text embeddings, e.g. <woman1>.
But when it goes to show cross-attention process:

prompts = [prompt]
controller = AttentionStore()
image_inv, x_t = run_and_display(prompts, controller, run_baseline=True)

It loses the information of <woman1>.
However, when I test with this:

image = ldm_stable(
      'A photo of <woman1>',
      num_inference_steps=50,
      guidance_scale=6.0,
      eta=1.0,
  ).images[0]

It returns the learned information of <woman1>.
How does this happened?