Add x-prediction / prediction_type=sample support for SDXL fine-tuning (#5095)

Co-authored-by: bghira <[email protected]>
Co-authored-by: Sayak Paul <[email protected]>

Author: 3 people

examples/text_to_image/train_text_to_image_sdxl.py

@@ -998,6 +998,11 @@ def compute_time_ids(original_size, crops_coords_top_left):
                     target = noise
                 elif noise_scheduler.config.prediction_type == "v_prediction":
                     target = noise_scheduler.get_velocity(model_input, noise, timesteps)
+                elif noise_scheduler.config.prediction_type == "sample":
+                    # We set the target to latents here, but the model_pred will return the noise sample prediction.
+                    target = model_input
+                    # We will have to subtract the noise residual from the prediction to get the target sample.
+                    model_pred = model_pred - noise
                 else:
                     raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")