Example of fine tuning an image-to-image model with Diffusers (StableImg2ImgPipeline)

[roomeo]

I want to kown how to design loss fuction:
def get_loss(data):
device = data['audio_embeds'].device

out_audio_embeds = data['audio_embeds']

prompt_embeds, negative_prompt_embeds = pipe.encode_prompt(
    prompt="",
    device=device,
    num_images_per_prompt=1,
    do_classifier_free_guidance=True,
    negative_prompt=None,
    prompt_embeds=None,
    negative_prompt_embeds=None,
    lora_scale=None,
)
prompt_embeds = prompt_embeds.to(out_audio_embeds.dtype)

image_embeds = encode_image(
    image=None,
    device=device,
    batch_size=1,
    num_images_per_prompt=1,
    do_classifier_free_guidance=True,
    noise_level=0,
    generator=None,
    image_embeds=out_audio_embeds,
)
image_embeds = image_embeds.to(out_audio_embeds.dtype)
pixel_values_embeds = data['pixel_values'].unsqueeze(dim=0)
out_vae = vae.encode(pixel_values_embeds).latent_dist.sample().detach()
out_vae = out_vae * 0.18215
noise = torch.randn_like(out_vae)
noise_step = torch.randint(0, 1000, (1,), device=device).long()
out_vae_noise = scheduler.add_noise(out_vae, noise, noise_step)
out_unet = unet(out_vae_noise, noise_step, encoder_hidden_states=prompt_embeds, class_labels=image_embeds).sample
return criterion(out_unet, noise)