thuanz123 / swiftbrush Goto Github PK

unet is loaded from pretrained sd1.5. 

tea_unet = copy.deepcopy(unet)
tea_unet.eval()
tea_unet.requires_grad_(False)

fake_unet = copy.deepcopy(unet)
fake_unet.train()
fake_unet.requires_grad_(True)

noise = torch.randn_like(latents)
bsz = noise.size(0)

with accelerator.accumulate(unet):

     noisy_latents = noise
     timesteps = torch.ones([bsz]).long().to(device) * (ddim_scheduler.num_train_timesteps - 1)
    
    pred_cond = unet(
        sample=noisy_latents,
        timestep=timesteps,
        encoder_hidden_states=text_positive_embeddings,
        attention_mask=None,
        added_cond_kwargs=added_cond_kwargs,
        ).sample

    predict = ddim_scheduler.convert_output(
        model_output=pred_cond,
        model_output_type=args.unet_prediction,
        sample=noisy_latents, 
        timesteps=timesteps,
        )
    predict_x0 = predict.pred_original_sample 
    predict_noise = predict.pred_epsilon

    ###################################
    # update fake score function    #
    ###################################
    # denoising Loss
    fake_timesteps = torch.randint(0, ddim_scheduler.num_train_timesteps, (bsz,), device=latents.device)
    fake_noise = torch.randn_like(predict_x0)

    fake_noisy_latents = ddim_scheduler.add_noise(predict_x0.detach(), fake_noise, fake_timesteps)

    fake_pred_cond = fake_unet(
        sample=fake_noisy_latents,
        timestep=fake_timesteps,
        encoder_hidden_states=text_positive_embeddings,
        attention_mask=None,
        added_cond_kwargs=added_cond_kwargs,
        ).sample
    fake_loss = F.mse_loss(fake_pred_cond.float(), fake_noise.float(), reduction="none")
    fake_snr = compute_snr(ddim_scheduler, fake_timesteps)
    fake_loss = (fake_loss.mean([1,2,3]) * fake_snr).mean()

    if (step + 1) % args.gradient_accumulation_steps == 0:  # debug check
        assert optimizer.gradient_state.sync_gradients is True
        assert optimizer_D.gradient_state.sync_gradients is True
    else:
        assert optimizer.gradient_state.sync_gradients is False
        assert optimizer_D.gradient_state.sync_gradients is False
    if accelerator.sync_gradients:
        accelerator.backward(fake_loss) 
        accelerator.clip_grad_norm_(fake_unet.parameters(), args.max_grad_norm)
    else:
        with fake_unet.no_sync():
            accelerator.backward(fake_loss)
            accelerator.clip_grad_norm_(fake_unet.parameters(), args.max_grad_norm)

    optimizer_D.step()
    optimizer_D.zero_grad()
    
    # distributionMatchingLoss
    dm_timesteps = torch.randint(int(ddim_scheduler.num_train_timesteps * 0.02), 
        int(ddim_scheduler.num_train_timesteps * 0.98), (bsz,), device=latents.device)
    dm_noise = torch.randn_like(predict_x0)
    dm_noisy_latents = ddim_scheduler.add_noise(predict_x0, dm_noise, timesteps)
    
    with torch.no_grad():
        pred_cond_real = tea_unet(
            sample=dm_noisy_latents,
            timestep=dm_timesteps,
            encoder_hidden_states=text_positive_embeddings,
            attention_mask=None,
            added_cond_kwargs=added_cond_kwargs,
            ).sample
        pred_uncond_real = tea_unet(
            sample=dm_noisy_latents,
            timestep=dm_timesteps,
            encoder_hidden_states=text_negative_embeddings,
            attention_mask=None,
            added_cond_kwargs=added_uncond_kwargs,
            ).sample
        
        pred_epsilon_real = classifier_free_guidance(
            pred_cond_real, pred_uncond_real, guidance_weight=args.validation_guidance_weight)

        pred_cond_fake = fake_unet(
            sample=dm_noisy_latents,
            timestep=dm_timesteps,
            encoder_hidden_states=text_positive_embeddings,
            attention_mask=None,
            added_cond_kwargs=added_cond_kwargs,
            ).sample
        pred_uncond_fake = fake_unet(
            sample=dm_noisy_latents,
            timestep=dm_timesteps,
            encoder_hidden_states=text_negative_embeddings,
            attention_mask=None,
            added_cond_kwargs=added_uncond_kwargs,
            ).sample
        pred_epsilon_fake = classifier_free_guidance(
            pred_cond_fake, pred_uncond_fake, guidance_weight=args.validation_guidance_weight)

    snr = compute_snr(ddim_scheduler, dm_timesteps)
    grad = (pred_epsilon_real - pred_epsilon_fake) * snr.reshape((bsz, 1, 1, 1))
    dm_loss = F.mse_loss(predict_x0, (predict_x0 - grad).detach()).mean()

    loss = dm_loss
    
    if accelerator.sync_gradients:
        accelerator.backward(loss)
        accelerator.clip_grad_norm_(unet.parameters(), args.max_grad_norm)
    else:
        with unet.no_sync():
            accelerator.backward(loss)
            accelerator.clip_grad_norm_(unet.parameters(), args.max_grad_norm)

    optimizer.step()
    optimizer.zero_grad()