Merge branch 'main' into dreambooth-lora-flux-exploration

Author: linoytsaban

benchmarks/push_results.py

@@ -3,7 +3,7 @@
 
 import pandas as pd
 from huggingface_hub import hf_hub_download, upload_file
-from huggingface_hub.utils._errors import EntryNotFoundError
+from huggingface_hub.utils import EntryNotFoundError
 
 
 sys.path.append(".")

src/diffusers/pipelines/controlnet_hunyuandit/pipeline_hunyuandit_controlnet.py

@@ -225,6 +225,8 @@ def __init__(
         requires_safety_checker: bool = True,
     ):
         super().__init__()
+        if isinstance(controlnet, (list, tuple)):
+            controlnet = HunyuanDiT2DMultiControlNetModel(controlnet)
 
         self.register_modules(
             vae=vae,

src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py

@@ -192,6 +192,8 @@ def __init__(
         ],
     ):
         super().__init__()
+        if isinstance(controlnet, (list, tuple)):
+            controlnet = SD3MultiControlNetModel(controlnet)
 
         self.register_modules(
             vae=vae,

src/diffusers/schedulers/scheduling_ddim.py

@@ -463,7 +463,10 @@ def step(
             prev_sample = prev_sample + variance
 
         if not return_dict:
-            return (prev_sample,)
+            return (
+                prev_sample,
+                pred_original_sample,
+            )
 
         return DDIMSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample)
 

src/diffusers/schedulers/scheduling_ddim_cogvideox.py

@@ -394,7 +394,10 @@ def step(
         prev_sample = a_t * sample + b_t * pred_original_sample
 
         if not return_dict:
-            return (prev_sample,)
+            return (
+                prev_sample,
+                pred_original_sample,
+            )
 
         return DDIMSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample)
 

src/diffusers/schedulers/scheduling_ddim_parallel.py

@@ -480,7 +480,10 @@ def step(
             prev_sample = prev_sample + variance
 
         if not return_dict:
-            return (prev_sample,)
+            return (
+                prev_sample,
+                pred_original_sample,
+            )
 
         return DDIMParallelSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample)
 

src/diffusers/schedulers/scheduling_ddpm.py

@@ -492,7 +492,10 @@ def step(
         pred_prev_sample = pred_prev_sample + variance
 
         if not return_dict:
-            return (pred_prev_sample,)
+            return (
+                pred_prev_sample,
+                pred_original_sample,
+            )
 
         return DDPMSchedulerOutput(prev_sample=pred_prev_sample, pred_original_sample=pred_original_sample)
 

src/diffusers/schedulers/scheduling_ddpm_parallel.py

@@ -500,7 +500,10 @@ def step(
         pred_prev_sample = pred_prev_sample + variance
 
         if not return_dict:
-            return (pred_prev_sample,)
+            return (
+                pred_prev_sample,
+                pred_original_sample,
+            )
 
         return DDPMParallelSchedulerOutput(prev_sample=pred_prev_sample, pred_original_sample=pred_original_sample)
 

src/diffusers/schedulers/scheduling_dpmsolver_sde.py

@@ -13,21 +13,41 @@
 # limitations under the License.
 
 import math
+from dataclasses import dataclass
 from typing import List, Optional, Tuple, Union
 
 import numpy as np
 import torch
 import torchsde
 
 from ..configuration_utils import ConfigMixin, register_to_config
-from ..utils import is_scipy_available
-from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput
+from ..utils import BaseOutput, is_scipy_available
+from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin
 
 
 if is_scipy_available():
     import scipy.stats
 
 
+@dataclass
+# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->DPMSolverSDE
+class DPMSolverSDESchedulerOutput(BaseOutput):
+    """
+    Output class for the scheduler's `step` function output.
+
+    Args:
+        prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images):
+            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
+            denoising loop.
+        pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images):
+            The predicted denoised sample `(x_{0})` based on the model output from the current timestep.
+            `pred_original_sample` can be used to preview progress or for guidance.
+    """
+
+    prev_sample: torch.Tensor
+    pred_original_sample: Optional[torch.Tensor] = None
+
+
 class BatchedBrownianTree:
     """A wrapper around torchsde.BrownianTree that enables batches of entropy."""
 
@@ -510,7 +530,7 @@ def step(
         sample: Union[torch.Tensor, np.ndarray],
         return_dict: bool = True,
         s_noise: float = 1.0,
-    ) -> Union[SchedulerOutput, Tuple]:
+    ) -> Union[DPMSolverSDESchedulerOutput, Tuple]:
         """
         Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
         process from the learned model outputs (most often the predicted noise).
@@ -522,15 +542,16 @@ def step(
                 The current discrete timestep in the diffusion chain.
             sample (`torch.Tensor` or `np.ndarray`):
                 A current instance of a sample created by the diffusion process.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or tuple.
+            return_dict (`bool`):
+                Whether or not to return a [`~schedulers.scheduling_dpmsolver_sde.DPMSolverSDESchedulerOutput`] or
+                tuple.
             s_noise (`float`, *optional*, defaults to 1.0):
                 Scaling factor for noise added to the sample.
 
         Returns:
-            [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`:
-                If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a
-                tuple is returned where the first element is the sample tensor.
+            [`~schedulers.scheduling_dpmsolver_sde.DPMSolverSDESchedulerOutput`] or `tuple`:
+                If return_dict is `True`, [`~schedulers.scheduling_dpmsolver_sde.DPMSolverSDESchedulerOutput`] is
+                returned, otherwise a tuple is returned where the first element is the sample tensor.
         """
         if self.step_index is None:
             self._init_step_index(timestep)
@@ -610,9 +631,12 @@ def t_fn(_sigma: torch.Tensor) -> torch.Tensor:
         self._step_index += 1
 
         if not return_dict:
-            return (prev_sample,)
+            return (
+                prev_sample,
+                pred_original_sample,
+            )
 
-        return SchedulerOutput(prev_sample=prev_sample)
+        return DPMSolverSDESchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample)
 
     # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.add_noise
     def add_noise(

src/diffusers/schedulers/scheduling_edm_euler.py

@@ -333,14 +333,13 @@ def step(
 
         gamma = min(s_churn / (len(self.sigmas) - 1), 2**0.5 - 1) if s_tmin <= sigma <= s_tmax else 0.0
 
-        noise = randn_tensor(
-            model_output.shape, dtype=model_output.dtype, device=model_output.device, generator=generator
-        )
-
-        eps = noise * s_noise
         sigma_hat = sigma * (gamma + 1)
 
         if gamma > 0:
+            noise = randn_tensor(
+                model_output.shape, dtype=model_output.dtype, device=model_output.device, generator=generator
+            )
+            eps = noise * s_noise
             sample = sample + eps * (sigma_hat**2 - sigma**2) ** 0.5
 
         # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
@@ -360,7 +359,10 @@ def step(
         self._step_index += 1
 
         if not return_dict:
-            return (prev_sample,)
+            return (
+                prev_sample,
+                pred_original_sample,
+            )
 
         return EDMEulerSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample)