Why does train_text_to_image.py perform so differently from the CompVis script?

[john-sungjin]

I posted about this on the forum but didn't get any useful feedback - would love to hear from someone who knows the in and outs of the diffusers codebase!

https://discuss.huggingface.co/t/discrepancies-between-compvis-and-diffuser-fine-tuning/25556

To summarize the post: the train_text_to_image.py script and original CompVis repo perform very differently when fine-tuning on the same dataset with the same hyperparameters. I'm trying to reproduce the Lamda Labs Pokemon fine-tuning results and finding difficulty doing so (picture results in forum post).

I've been digging into the implementations and I'm not noticing any obvious differences in how the models are trained, losses are calculated, etc - so what explains the large behavioral discrepancies?

Would really appreciate any insight on what might be causing this.

[patrickvonplaten]

Thanks a lot for posting the question here @john-sungjin!

Sorry for the moment we try to keep everything on GitHub. From the discussion it seems like many people are encountering your issue so let's maybe spend some time to figure out what's going on. @patil-suraj do you think you might have some time soon to look into it? If not I can try to allocate time (or maybe cc @anton-l)

[patil-suraj]

Thanks for posting the detailed issue @john-sungjin !

As you said, the implementation is very similar to the compvis one. The one difference that I'm aware of is that, in the compvis script, for example the Pokemon fine-tuning script, the model is initialised from the sd-v1-4-full-ema.ckpt checkpoint, so it loads the non-ema weights for training and ema weights for doing ema. While in diffusers script the ema checkpoint is used for both training and EMA.

I am going to add an option which enables loading both the non-ema (for training) and ema (for EMA updates) in diffusers script and then compare again. Will report here as soon as possible :)

[AIXiaoBaiDemon]

I had the same problem, looking forward to your experiment。

[treksis]

same problem

[Line290]

Mark it

[patil-suraj]

Going to update the script soon, I am getting good results with script now, see for example the emoji model

[Line290]

Hi @patil-suraj ，
I tried to load unet_with_ema in EMAModel and unet_without_ema for fine-tuning, but the result remains terrible. Could you help me fix it? Thanks.

Here are my code snippets:

class EMAModel:
    """
    Exponential Moving Average of models weights
    """

    def __init__(self, parameters: Iterable[torch.nn.Parameter], device, dtype, decay=0.9999):
        parameters = list(parameters)
        self.shadow_params = [p.clone().detach().to(device=device, dtype=dtype)
                              for p in parameters]
        self.decay = decay
        self.optimization_step = 0

unet_ema = UNet2DConditionModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="unet")
unet = UNet2DConditionModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="unet_wo_ema")

# Create EMA for the unet.
if args.use_ema:
    ema_unet = EMAModel(unet_ema.parameters(), device=accelerator.device, dtype=weight_dtype)

[haofanwang]

Any update? The results are bad as least on the pokemon dataset.

[patrickvonplaten]

Ping @patil-suraj here

[patrickvonplaten]

@patil-suraj can you please make the required updates? Also cc @williamberman

[haofanwang]

The problem goes away in the lastest 0.10.0 version. But it would be helpful to show what specific modifications make it work.

[patrickvonplaten]

Another ping here @patil-suraj - could you look into this. Should we show how to use the EMA weights?