Questions about supporting v-prediction in SD21 models

[loggerJK]

Hi,

First, thank you for your fantastic work on this project! I have a couple of questions regarding the implementation of SD21 in the codebase:

(1) It seems the current codebase includes the SD21 model, but it doesn't fully function as expected due to the implementation only supporting epsilon prediction. Is there any way to add support for v-prediction for SD21?
(2) Is there a theoretical reason behind only implementing epsilon prediction, or is it more of a practical limitation?

Thank you for your time and help!

[lxueaa]

I tried the following settings, it works for SD2.1:
Firstly, I deal with the v-prediction problem. I delete the line "prediction_type": "v_prediction" in runwayml/stable-diffusion-v1-5/scheduler/scheduler_config.json, forcing DDIM to use epsilon. Then the code is:

import torch
from PIL import Image
import matplotlib.pyplot as plt

from src.eunms import Model_Type, Scheduler_Type
from src.utils.enums_utils import get_pipes
from src.config import RunConfig
from main import run as invert

device = "cuda:0" if torch.cuda.is_available() else "cpu"
model_type = Model_Type.SD21
scheduler_type = Scheduler_Type.DDIM

pipe_inversion, pipe_inference = get_pipes(
    model_type=model_type,
    scheduler_type=scheduler_type,
    device=device,
)

cfg = RunConfig(
    model_type=model_type,
    scheduler_type=scheduler_type,
    # you can tweak these:
    num_inference_steps=50,      # denoising steps 
    num_inversion_steps=50,      # inversion steps
    guidance_scale=0.0,         # classifier-free guidance during inversion
    num_renoise_steps=7,        # ReNoise iterations per step
    max_num_renoise_steps_first_step=5,
    inversion_max_step=1.0,     # 1.0 = full noise, 0.5 = half steps, etc.
    average_latent_estimations=True,
    average_first_step_range=(0, 5),
    average_step_range=(8, 10),
    noise_regularization_lambda_ac=10.0,
    noise_regularization_lambda_kl=0.055,
    noise_regularization_num_reg_steps=4,
    noise_regularization_num_ac_rolls=5,
    perform_noise_correction=True,
    seed=7865,
)

img_path = "example_images/kitten.jpg"
input_image = (
    Image.open(img_path)
    .convert("RGB")
    .resize((1024, 1024))
)

prompt = "a cat in the photo"
rec_img, inv_latent, noise, all_latents = invert(
    init_image=input_image,
    prompt=prompt,
    cfg=cfg,
    pipe_inversion=pipe_inversion,
    pipe_inference=pipe_inference,
    do_reconstruction=True,
)

fig, axes = plt.subplots(1, 2, figsize=(10, 3))

axes[0].imshow(input_image)
axes[0].set_title("Origin")
axes[0].axis("off")

axes[1].imshow(rec_img)
axes[1].set_title("Reconstruction")
axes[1].axis("off")

Intuitively, I think v-prediction is not stable compared with epsilon.