o ۷iec@sddlZddlmZddlmZddlZddlmZmZddl m Z m Z ddl m Z dd lmZe eZeGd d d e ZGd d d eeZdS)N) dataclass)Literal) ConfigMixinregister_to_config) BaseOutputlogging) randn_tensor)SchedulerMixinc@s.eZdZUdZejed<dZejdBed<dS)EDMEulerSchedulerOutputaq 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_sampleNpred_original_sample)__name__ __module__ __qualname____doc__torchTensor__annotations__rrr_/home/ubuntu/vllm_env/lib/python3.10/site-packages/diffusers/schedulers/scheduling_edm_euler.pyr s  r c@seZdZdZgZdZe       dUd ed ed ededde deddededddfddZ e defddZ e de fddZ e de fddZdVd!e ddfd"d#Zd$ejd%eejBdejfd&d'Zd%eejBdejfd(d)Zd$ejd*ejd%eejBdejfd+d,Zd$ejd-eejBdejfd.d/Z   dWd0e d1eejBd2ejeeBdBfd3d4Z  dXd5ejd edBd edBdejfd6d7Z  dXd5ejd edBd edBdejfd8d9Z dYd-eejBd:ejdBde fd;d<Zd-eejBddfd=d>Zd?d?ed@dAddBdfd*ejd-eejBd$ejdCedDedEedFedGejdBdHedIejdBde e!BfdJdKZ"dLejdMejdNejdejfdOdPZ#d%eejBdeejBfdQdRZ$de fdSdTZ%dS)ZEDMEulerScheduleraL Implements the Euler scheduler in EDM formulation as presented in Karras et al. 2022 [1]. [1] Karras, Tero, et al. "Elucidating the Design Space of Diffusion-Based Generative Models." https://huggingface.co/papers/2206.00364 This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic methods the library implements for all schedulers such as loading and saving. Args: sigma_min (`float`, *optional*, defaults to `0.002`): Minimum noise magnitude in the sigma schedule. This was set to 0.002 in the EDM paper [1]; a reasonable range is [0, 10]. sigma_max (`float`, *optional*, defaults to `80.0`): Maximum noise magnitude in the sigma schedule. This was set to 80.0 in the EDM paper [1]; a reasonable range is [0.2, 80.0]. sigma_data (`float`, *optional*, defaults to `0.5`): The standard deviation of the data distribution. This is set to 0.5 in the EDM paper [1]. sigma_schedule (`Literal["karras", "exponential"]`, *optional*, defaults to `"karras"`): Sigma schedule to compute the `sigmas`. By default, we use the schedule introduced in the EDM paper (https://huggingface.co/papers/2206.00364). The `"exponential"` schedule was incorporated in this model: https://huggingface.co/stabilityai/cosxl. num_train_timesteps (`int`, *optional*, defaults to `1000`): The number of diffusion steps to train the model. prediction_type (`Literal["epsilon", "v_prediction"]`, *optional*, defaults to `"epsilon"`): Prediction type of the scheduler function. `"epsilon"` predicts the noise of the diffusion process, and `"v_prediction"` (see section 2.4 of [Imagen Video](https://huggingface.co/papers/2210.02303) paper). rho (`float`, *optional*, defaults to `7.0`): The rho parameter used for calculating the Karras sigma schedule, which is set to 7.0 in the EDM paper [1]. final_sigmas_type (`Literal["zero", "sigma_min"]`, *optional*, defaults to `"zero"`): The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final sigma is the same as the last sigma in the training schedule. If `"zero"`, the final sigma is set to 0. r Mb`?T@?karrasepsilon@zero sigma_min sigma_max sigma_datasigma_scheduler exponentialnum_train_timestepsprediction_type)r v_predictionrhofinal_sigmas_type)r r!returnNc Cs|dvr td|dd|_tjjrtjntj} tj|d| d|} |dkr0| | } n |dkr9| | } | tj} | | |_ |jjdkrP| d } n|jjd krYd } n td |jjt| tjd | | jdg|_d|_d|_d|_|j d|_dS)Nr%z-Wrong value for provided for `sigma_schedule=z`.`r dtyperr&r!r rC`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got r  fill_valuedeviceFcpu) ValueErrornum_inference_stepsrbackendsmps is_availablefloat32float64arange_compute_karras_sigmas_compute_exponential_sigmastoprecondition_noise timestepsconfigr+catfullr4sigmasis_scale_input_called _step_index _begin_index) selfr!r"r#r$r'r(r*r+ sigmas_dtyperF sigma_lastrrr__init__Ws.         zEDMEulerScheduler.__init__cCs|jjdddS)z Return the standard deviation of the initial noise distribution. Returns: `float`: The initial noise sigma value computed as `(sigma_max**2 + 1) ** 0.5`. rr r)rCr"rJrrrinit_noise_sigmas z"EDMEulerScheduler.init_noise_sigmacC|jS)z Return the index counter for the current timestep. The index will increase by 1 after each scheduler step. Returns: `int` or `None`: The current step index, or `None` if not yet initialized. )rHrNrrr step_indexs zEDMEulerScheduler.step_indexcCrP)z Return the index for the first timestep. This should be set from the pipeline with the `set_begin_index` method. Returns: `int` or `None`: The begin index, or `None` if not yet set. rIrNrrr begin_indexs zEDMEulerScheduler.begin_indexrrScCs ||_dS)z Sets the begin index for the scheduler. This function should be run from pipeline before the inference. Args: begin_index (`int`, defaults to `0`): The begin index for the scheduler. NrR)rJrSrrrset_begin_indexs z!EDMEulerScheduler.set_begin_indexsamplesigmacCs||}||}|S)a Precondition the input sample by scaling it according to the EDM formulation. Args: sample (`torch.Tensor`): The input sample tensor to precondition. sigma (`float` or `torch.Tensor`): The current sigma (noise level) value. Returns: `torch.Tensor`: The scaled input sample. )_get_conditioning_c_in)rJrUrVc_in scaled_samplerrrprecondition_inputss z%EDMEulerScheduler.precondition_inputscCs*t|tjs t|g}dt|}|S)aS Precondition the noise level by applying a logarithmic transformation. Args: sigma (`float` or `torch.Tensor`): The sigma (noise level) value to precondition. Returns: `torch.Tensor`: The preconditioned noise value computed as `0.25 * log(sigma)`. g?) isinstancerrtensorlog)rJrVc_noiserrrrAs z$EDMEulerScheduler.precondition_noise model_outputcCs|jj}|d|d|d}|jjdkr%|||d|dd}n |jjdkr;| ||d|dd}n td|jjd||||}|S)a Precondition the model outputs according to the EDM formulation. Args: sample (`torch.Tensor`): The input sample tensor. model_output (`torch.Tensor`): The direct output from the learned diffusion model. sigma (`float` or `torch.Tensor`): The current sigma (noise level) value. Returns: `torch.Tensor`: The denoised sample computed by combining the skip connection and output scaling. rrrr)zPrediction type z is not supported.)rCr#r(r6)rJrUr_rVr#c_skipc_outdenoisedrrrprecondition_outputss   z&EDMEulerScheduler.precondition_outputstimestepcCs6|jdur |||j|j}|||}d|_|S)a Scale the denoising model input to match the Euler algorithm. Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep. Args: sample (`torch.Tensor`): The input sample tensor. timestep (`float` or `torch.Tensor`): The current timestep in the diffusion chain. Returns: `torch.Tensor`: A scaled input sample. NT)rQ_init_step_indexrFrZrG)rJrUrdrVrrrscale_model_inputs    z#EDMEulerScheduler.scale_model_inputr7r4rFcCs ||_tjjr tjntj}|durtjdd|j|d}nt|t r+tj ||d}n| |}|j j dkr<||}n |j j dkrG||}|j tj|d}|||_|j jdkr`|d }n|j jd krid}n td |j jt|tjd ||jd g|_d|_d|_|j d|_dS)a Sets the discrete timesteps used for the diffusion chain (to be run before inference). Args: num_inference_steps (`int`, *optional*): The number of diffusion steps used when generating samples with a pre-trained model. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. sigmas (`torch.Tensor | list[float]`, *optional*): Custom sigmas to use for the denoising process. If not defined, the default behavior when `num_inference_steps` is passed will be used. Nrr r-rr&)r.r4r!r/r r0r1r2r5)r7rr8r9r:r;r<linspacer[floatr\r@rCr$r>r?rArBr+r6rDrEr4rFrHrI)rJr7r4rFrKrLrrr set_timestepss0           zEDMEulerScheduler.set_timestepsrampcCsP|p|jj}|p |jj}|jj}|d|}|d|}|||||}|S)aT Construct the noise schedule of [Karras et al. (2022)](https://huggingface.co/papers/2206.00364). Args: ramp (`torch.Tensor`): A tensor of values in [0, 1] representing the interpolation positions. sigma_min (`float`, *optional*): Minimum sigma value. If `None`, uses `self.config.sigma_min`. sigma_max (`float`, *optional*): Maximum sigma value. If `None`, uses `self.config.sigma_max`. Returns: `torch.Tensor`: The computed Karras sigma schedule. r )rCr!r"r*)rJrjr!r"r* min_inv_rho max_inv_rhorFrrrr>Cs    z(EDMEulerScheduler._compute_karras_sigmascCsD|p|jj}|p |jj}tt|t|t| d}|S)a Compute the exponential sigma schedule. Implementation closely follows k-diffusion: https://github.com/crowsonkb/k-diffusion/blob/6ab5146d4a5ef63901326489f31f1d8e7dd36b48/k_diffusion/sampling.py#L26 Args: ramp (`torch.Tensor`): A tensor of values representing the interpolation positions. sigma_min (`float`, *optional*): Minimum sigma value. If `None`, uses `self.config.sigma_min`. sigma_max (`float`, *optional*): Maximum sigma value. If `None`, uses `self.config.sigma_max`. Returns: `torch.Tensor`: The computed exponential sigma schedule. r) rCr!r"rrgmathr]lenexpflip)rJrjr!r"rFrrrr?as  (z-EDMEulerScheduler._compute_exponential_sigmasschedule_timestepscCs:|dur|j}||k}t|dkrdnd}||S)ak Find the index of a given timestep in the timestep schedule. Args: timestep (`float` or `torch.Tensor`): The timestep value to find in the schedule. schedule_timesteps (`torch.Tensor`, *optional*): The timestep schedule to search in. If `None`, uses `self.timesteps`. Returns: `int`: The index of the timestep in the schedule. For the very first step, returns the second index if multiple matches exist to avoid skipping a sigma when starting mid-schedule (e.g., for image-to-image). Nr r)rBnonzerornitem)rJrdrqindicesposrrrindex_for_timestep}s   z$EDMEulerScheduler.index_for_timestepcCs@|jdurt|tjr||jj}|||_dS|j |_dS)z Initialize the step index for the scheduler based on the given timestep. Args: timestep (`float` or `torch.Tensor`): The current timestep to initialize the step index from. N) rSr[rrr@rBr4rvrHrI)rJrdrrrres   z"EDMEulerScheduler._init_step_indexinfg?Ts_churns_tmins_tmaxs_noise generator return_dictrc CsNt|ttjtjfrtd|jstd|j dur | || tj }|j |j } || kr6|krDnn t|t|j ddnd} | | d} | dkrmt|j|j|j|d}||}||| d | d d }| durx|||| } || | }|j |j d| }|||}| |j}|jd7_| s|| fSt|| d S) a 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). Args: model_output (`torch.Tensor`): The direct output from the learned diffusion model. timestep (`float` or `torch.Tensor`): The current discrete timestep in the diffusion chain. sample (`torch.Tensor`): A current instance of a sample created by the diffusion process. s_churn (`float`, *optional*, defaults to `0.0`): The amount of stochasticity to add at each step. Higher values add more noise. s_tmin (`float`, *optional*, defaults to `0.0`): The minimum sigma threshold below which no noise is added. s_tmax (`float`, *optional*, defaults to `float("inf")`): The maximum sigma threshold above which no noise is added. s_noise (`float`, *optional*, defaults to `1.0`): Scaling factor for noise added to the sample. generator (`torch.Generator`, *optional*): A random number generator for reproducibility. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return an [`~schedulers.scheduling_edm_euler.EDMEulerSchedulerOutput`] or tuple. pred_original_sample (`torch.Tensor`, *optional*): The predicted denoised sample from a previous step. If provided, skips recomputation. Returns: [`~schedulers.scheduling_edm_euler.EDMEulerSchedulerOutput`] or `tuple`: If `return_dict` is `True`, an [`~schedulers.scheduling_edm_euler.EDMEulerSchedulerOutput`] is returned, otherwise a tuple is returned where the first element is the previous sample tensor and the second element is the predicted original sample tensor. zPassing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to `EDMEulerScheduler.step()` is not supported. Make sure to pass one of the `scheduler.timesteps` as a timestep.zThe `scale_model_input` function should be called before `step` to ensure correct denoising. See `StableDiffusionPipeline` for a usage example.Nr g4y?rwr)r.r4r}rr)r r)r[intr IntTensor LongTensorr6rGloggerwarningrQrer@r;rFminrnr shaper.r4rcrHr )rJr_rdrUryrzr{r|r}r~rrVgamma sigma_hatnoiseeps derivativedtr rrrstepsF.    4     zEDMEulerScheduler.steporiginal_samplesrrBcs jj|j|jd}|jjdkr)t|r)jj|jtjd|j|jtjd}n j|j||j}j durFfdd|D}nj durUj g|j d}n j g|j d}|| }t |j t |j kr}|d}t |j t |j ksn|||}|S) am Add noise to the original samples according to the noise schedule at the specified timesteps. Args: original_samples (`torch.Tensor`): The original samples to which noise will be added. noise (`torch.Tensor`): The noise tensor to add to the original samples. timesteps (`torch.Tensor`): The timesteps at which to add noise, determining the noise level from the schedule. Returns: `torch.Tensor`: The noisy samples with added noise scaled according to the timestep schedule. )r4r.r9r-Ncsg|]}|qSr)rv).0trqrJrr 9sz/EDMEulerScheduler.add_noise..rr/)rFr@r4r.typeris_floating_pointrBr;rSrQrflattenrn unsqueeze)rJrrrBrF step_indicesrV noisy_samplesrrr add_noises"      zEDMEulerScheduler.add_noisecCs d|d|jjdd}|S)a4 Compute the input conditioning factor for the EDM formulation. Args: sigma (`float` or `torch.Tensor`): The current sigma (noise level) value. Returns: `float` or `torch.Tensor`: The input conditioning factor `c_in`. r rr)rCr#)rJrVrXrrrrWHs z(EDMEulerScheduler._get_conditioning_c_incCs|jjSN)rCr'rNrrr__len__WszEDMEulerScheduler.__len__)rrrrrrrr )r)NNN)NNr)&rrrr _compatiblesorderrrhrrrMpropertyrOrQrSrTrrrZrArcrfstrr4listrir>r?rvre Generatorboolr tuplerrrWrrrrrr1s "  ,    # 5 !      m 0r)rm dataclassesrtypingrrconfiguration_utilsrrutilsrrutils.torch_utilsr scheduling_utilsr get_loggerrrr rrrrrs