EMA: fix state_dict() and load_state_dict() & add cur_decay_value

[chenguolin]

1. fix the saved value for min_decay in EMA self.state_dict().

2. add an interface self.cur_decay_value to track the current value for decay.
   (as self.decay is a constant value meaning "max decay value")

3. track the current EMA decay value by self.cur_decay_value in "unconditional image generation" examples, instead of self.decay.

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[chenguolin]

4. fix a bug in EMA load_static_dict(): "AttributeError: 'float' object has no attribute 'get'"

[patrickvonplaten]

This looks good to me! @chenguolin could you maybe also check if the EMAModel class is correctly used in: https://github.com/huggingface/diffusers/blob/main/examples/unconditional_image_generation/train_unconditional.py and potentially fix it? :-)

[chenguolin]

It looks good. I have changed logs["ema_decay"] to the new ema_model.cur_decay_value to track the current decay value at:

diffusers/examples/unconditional_image_generation/train_unconditional.py

Line 522 in 7d96b38

logs["ema_decay"] = ema_model.decay

[patrickvonplaten]

This LGTM, @pcuenca @patil-suraj could you take a look?

[patrickvonplaten]

Actually we should probably also apply the same changes to: https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py

In case you have 5min it'd be great if you could take a look at applying the same changes to train_text_to_image as wel @chenguolin :-)

cc @patil-suraj wdyt of this PR

[chenguolin]

Hi @patrickvonplaten, the only necessary change to example "train_unconditional.py" is logs["ema_decay"] = ema_model.decay -> logs["ema_decay"] = ema_model.cur_decay_value.

While other examples don't use ema_model.decay, maybe there is no need to change them.

[patil-suraj]

Thanks a lot for fixing this!

[patil-suraj]

@chenguolin there's a merge conflict in train_unconditional_ort.py could you please fix it? It should be good to merge after that.

[pcuenca]

Thank you!

[chenguolin]

Hi @patil-suraj, I have just deleted train_unconditional_ort.py ane merged.

[williamberman]

Failing test is un-related :)