KeyError: 'state_dict' not fix about sd_dreambooth_extension HOT 9 CLOSED

Just pushed f033e79 that should hopefully address this.

from sd_dreambooth_extension.

@d8ahazard
still have porblems create model with animefull-latest.ckpt [e6e8e1fc], but no problems with other models.

from sd_dreambooth_extension.

So, this one is tricky.

I can't find that checkpoint directly, so I can't look at why it's breaking, but it seems as if the previous fix in trying to load the checkpoint from a different path doesn't help. If you can provide a URL to the ckpt, I can look at the contents and try to determine why it's failing.

from sd_dreambooth_extension.

@d8ahazard
https://pub-2fdef7a2969f43289c42ac5ae3412fd4.r2.dev/animefull-latest.ckpt
fact this model was nice with training on dreambooth for colab. I hop this problem could be solved.

from sd_dreambooth_extension.

I think it is caused by that you have use a special ckpt model file which doesn't has a 'state_dict' key", and the follow code may help you:

  sd = torch.load(model_path, map_location="cpu")
  if "state_dict" not in sd:
      pruned_sd = {
          "state_dict": dict(),
      }
  else:
      pruned_sd = dict()
  for k in sd.keys():
      if k != "optimizer_states":
          if "state_dict" not in sd:
              pruned_sd["state_dict"][k] = sd[k]
          else:
              pruned_sd[k] = sd[k]
  torch.save(pruned_sd, "model-pruned.ckpt")

just remove "optimizer_states" key, and put other key to "state_dict" to get a normal model for you.
So that the code:

checkpoint = torch.load(checkpoint_file[0])['state_dict']

may work well.

from sd_dreambooth_extension.

I think it is caused by that you have use a special ckpt model file which doesn't has a 'state_dict' key", and the follow code may help you:

  sd = torch.load(model_path, map_location="cpu")
  if "state_dict" not in sd:
      pruned_sd = {
          "state_dict": dict(),
      }
  else:
      pruned_sd = dict()
  for k in sd.keys():
      if k != "optimizer_states":
          if "state_dict" not in sd:
              pruned_sd["state_dict"][k] = sd[k]
          else:
              pruned_sd[k] = sd[k]
  torch.save(pruned_sd, "model-pruned.ckpt")

just remove "optimizer_states" key, and put other key to "state_dict" to get a normal model for you. So that the code:

checkpoint = torch.load(checkpoint_file[0])['state_dict']

may work well.

Oh thanks,it works well

from sd_dreambooth_extension.

from sd_dreambooth_extension.

I'm getting the same error messages as 1301715759 user, except that I'm using sd-v1-4.ckpt, which is the default model. It happens when trying to Create a new model, so I'm not even to the point I can try Training yet. I'm stuck at the first step.
I have no idea how to implement the solution proposed by CrazyBoyM, and I know every time I try to change code from a file I'll get in troubles later (things like : "that file is no good you need to commit first blabla" and then I can't update properly again)

from sd_dreambooth_extension.

Closing, fixed with a9bbe70

from sd_dreambooth_extension.