o Gih@sddlZddlmZddlmZmZmZddlZddl Z ddl m Z m Z ddl mZmZmZddlmZer;ddlZeeZeGd d d eZGd d d ee ZdS) N) dataclass)LiteralOptionalUnion) ConfigMixinregister_to_config) BaseOutputis_scipy_availablelogging)SchedulerMixinc@seZdZUdZejed<dS)%FlowMatchEulerDiscreteSchedulerOutputaJ Output class for the scheduler's `step` function output. Args: prev_sample (`torch.FloatTensor` 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. prev_sampleN)__name__ __module__ __qualname____doc__torch FloatTensor__annotations__rrm/home/ubuntu/.local/lib/python3.10/site-packages/diffusers/schedulers/scheduling_flow_match_euler_discrete.pyr!s  rc@seZdZdZgZdZe           d`d ed ede ded Bded Bdedede dede de de de dde fddZ e ddZ e dd Ze d!d"Zdad$efd%d&Zd efd'd(Z dbd)ejd*eejBd+ejd Bd,ejfd-d.Zd,efd/d0Zd1ed2ed3ejd,ejfd4d5Zd3ejd,ejfd6d7Z dcd8ed Bd9eejBd:eed Bd1ed Bd;eed Bf deejd,efd?d@Zd*eeejfd,d fdAdBZ dCdCedDdd d dEfdFejd*eejBd)ejdGedHedIedJedKej!d BdLejd BdMe d,e"e#BfdNdOZ$dPejd8ed,ejfdQdRZ%dPejd8ed,ejfdSdTZ& UdddPejd8edVedWed,ejf dXdYZ'd1ed2ed3ejd,ejfdZd[Z(d1ed2ed3ejd,ejfd\d]Z)d,efd^d_Z*d S)eFlowMatchEulerDiscreteSchedulerab Euler scheduler. 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: num_train_timesteps (`int`, defaults to 1000): The number of diffusion steps to train the model. shift (`float`, defaults to 1.0): The shift value for the timestep schedule. use_dynamic_shifting (`bool`, defaults to False): Whether to apply timestep shifting on-the-fly based on the image resolution. base_shift (`float`, defaults to 0.5): Value to stabilize image generation. Increasing `base_shift` reduces variation and image is more consistent with desired output. max_shift (`float`, defaults to 1.15): Value change allowed to latent vectors. Increasing `max_shift` encourages more variation and image may be more exaggerated or stylized. base_image_seq_len (`int`, defaults to 256): The base image sequence length. max_image_seq_len (`int`, defaults to 4096): The maximum image sequence length. invert_sigmas (`bool`, defaults to False): Whether to invert the sigmas. shift_terminal (`float`, defaults to None): The end value of the shifted timestep schedule. use_karras_sigmas (`bool`, defaults to False): Whether to use Karras sigmas for step sizes in the noise schedule during sampling. use_exponential_sigmas (`bool`, defaults to False): Whether to use exponential sigmas for step sizes in the noise schedule during sampling. use_beta_sigmas (`bool`, defaults to False): Whether to use beta sigmas for step sizes in the noise schedule during sampling. time_shift_type (`str`, defaults to "exponential"): The type of dynamic resolution-dependent timestep shifting to apply. Either "exponential" or "linear". stochastic_sampling (`bool`, defaults to False): Whether to use stochastic sampling. r ?F?ffffff?N exponentialnum_train_timestepsshiftuse_dynamic_shifting base_shift max_shiftbase_image_seq_lenmax_image_seq_len invert_sigmasshift_terminaluse_karras_sigmasuse_exponential_sigmasuse_beta_sigmastime_shift_type)r linearstochastic_samplingcCs|jjr ts tdt|jj|jj|jjgdkrtd| dvr&tdtj d||tj dddd }t |jt j d}||}|sS||d|d|}|||_d|_d|_||_|d|_|jd|_|jd |_dS) Nz:Make sure to install scipy if you want to use beta sigmas.r znOnly one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used.>r.r z;`time_shift_type` must either be 'exponential' or 'linear'.dtypecpur)configr,r ImportErrorsumr+r* ValueErrornplinspacefloat32copyr from_numpyto timesteps _step_index _begin_index_shiftsigmasitem sigma_min sigma_max)selfr!r"r#r$r%r&r'r(r)r*r+r,r-r/r>rBrrr__init__Zs6 "  z(FlowMatchEulerDiscreteScheduler.__init__cC|jS)z. The value used for shifting. rArFrrrr"z%FlowMatchEulerDiscreteScheduler.shiftcCrH)zg The index counter for current timestep. It will increase 1 after each scheduler step. )r?rJrrr step_indexrKz*FlowMatchEulerDiscreteScheduler.step_indexcCrH)zq The index for the first timestep. It should be set from pipeline with `set_begin_index` method. r@rJrrr begin_indexrKz+FlowMatchEulerDiscreteScheduler.begin_indexrrNcC ||_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. NrM)rFrNrrrset_begin_index z/FlowMatchEulerDiscreteScheduler.set_begin_indexcCrO)z Sets the shift value for the scheduler. Args: shift (`float`): The shift value to be set. NrI)rFr"rrr set_shiftrQz)FlowMatchEulerDiscreteScheduler.set_shiftsampletimestepnoisereturncsjj|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||d ||}|S) a Forward process in flow-matching Args: sample (`torch.FloatTensor`): The input sample. timestep (`torch.FloatTensor`): The current timestep in the diffusion chain. noise (`torch.FloatTensor`): The noise tensor. Returns: `torch.FloatTensor`: A scaled input sample. )devicer1mpsr0Ncsg|]}|qSr)index_for_timestep).0tschedule_timestepsrFrr sz?FlowMatchEulerDiscreteScheduler.scale_noise..rr2r)rBr=rWr1typeris_floating_pointr>r:rNrLshapeflattenlen unsqueeze)rFrSrTrUrB step_indicessigmarr\r scale_noises"     z+FlowMatchEulerDiscreteScheduler.scale_noisecCs ||jjSNr4r!)rFrfrrr _sigma_to_ts z+FlowMatchEulerDiscreteScheduler._sigma_to_tmurfr[cCs8|jjdkr ||||S|jjdkr||||SdS)a Apply time shifting to the sigmas. Args: mu (`float`): The mu parameter for the time shift. sigma (`float`): The sigma parameter for the time shift. t (`torch.Tensor`): The input timesteps. Returns: `torch.Tensor`: The time-shifted timesteps. r r.N)r4r-_time_shift_exponential_time_shift_linearrFrkrfr[rrr time_shifts  z*FlowMatchEulerDiscreteScheduler.time_shiftcCs,d|}|dd|jj}d||}|S)a, Stretches and shifts the timestep schedule to ensure it terminates at the configured `shift_terminal` config value. Reference: https://github.com/Lightricks/LTX-Video/blob/a01a171f8fe3d99dce2728d60a73fecf4d4238ae/ltx_video/schedulers/rf.py#L51 Args: t (`torch.Tensor`): A tensor of timesteps to be stretched and shifted. Returns: `torch.Tensor`: A tensor of adjusted timesteps such that the final value equals `self.config.shift_terminal`. r r2)r4r))rFr[ one_minus_z scale_factor stretched_trrrstretch_shift_to_terminals z9FlowMatchEulerDiscreteScheduler.stretch_shift_to_terminalnum_inference_stepsrWrBr>cCsF|jjr |dur td|dur |dur t|t|kr td|dur=|dur.t||ks8|durrBr?r@)rFrtrWrBrkr>is_timesteps_providedrrr set_timestepssf      z-FlowMatchEulerDiscreteScheduler.set_timestepsr]cCs:|dur|j}||k}t|dkrdnd}||S)a Get the index for the given timestep. Args: timestep (`float` or `torch.FloatTensor`): The timestep to find the index for. schedule_timesteps (`torch.FloatTensor`, *optional*): The schedule timesteps to validate against. If `None`, the scheduler's timesteps are used. Returns: `int`: The index of the timestep. Nr r)r>nonzerorcrC)rFrTr]indicesposrrrrYs   z2FlowMatchEulerDiscreteScheduler.index_for_timestepcCs@|jdurt|tjr||jj}|||_dS|j |_dSrh) rN isinstancerTensorr=r>rWrYr?r@)rFrTrrr_init_step_indexs   z0FlowMatchEulerDiscreteScheduler._init_step_indexginfT model_outputs_churns_tmins_tmaxs_noise generatorper_token_timesteps return_dictc CsXt|tst|tjst|tjrtd|jdur|||tj }| durZ| |j j } |j ddddf} | | ddk} | | }|j dd\}}| d}|d}||}n|j}|j |}|j |d}|}|}||}|j jr|||}t|}d||||}n|||}|jd7_| dur||j}| 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.FloatTensor`): The direct output from learned diffusion model. timestep (`float`): The current discrete timestep in the diffusion chain. sample (`torch.FloatTensor`): A current instance of a sample created by the diffusion process. s_churn (`float`): s_tmin (`float`): s_tmax (`float`): s_noise (`float`, defaults to 1.0): Scaling factor for noise added to the sample. generator (`torch.Generator`, *optional*): A random number generator. per_token_timesteps (`torch.Tensor`, *optional*): The timesteps for each token in the sample. return_dict (`bool`, defaults to `True`): Whether or not to return a [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or tuple. Returns: [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or `tuple`: If return_dict is `True`, [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] is returned, otherwise a tuple is returned where the first element is the sample tensor. zPassing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to `FlowMatchEulerDiscreteScheduler.step()` is not supported. Make sure to pass one of the `scheduler.timesteps` as a timestep.Ngư>r)dim).Nr r)r)rintr IntTensor LongTensorr7rLrr=r:r4r!rBmaxr/ randn_liker?r1r)rFrrTrSrrrrrrrper_token_sigmasrB lower_mask lower_sigmas_ current_sigma next_sigmadt sigma_idxrf sigma_nextx0rUrrrrstepsL-             z$FlowMatchEulerDiscreteScheduler.stepruc Cst|jdr |jj}nd}t|jdr|jj}nd}|dur |n|d}|dur,|n|d}d}tdd|}|d|}|d|}|||||} | S)a Construct the noise schedule as proposed in [Elucidating the Design Space of Diffusion-Based Generative Models](https://huggingface.co/papers/2206.00364). Args: in_sigmas (`torch.Tensor`): The input sigma values to be converted. num_inference_steps (`int`): The number of inference steps to generate the noise schedule for. Returns: `torch.Tensor`: The converted sigma values following the Karras noise schedule. rDNrEr2rg@r )hasattrr4rDrErCr8r9) rFrurtrDrErhoramp min_inv_rho max_inv_rhorBrrrrxs      z2FlowMatchEulerDiscreteScheduler._convert_to_karrascCst|jdr |jj}nd}t|jdr|jj}nd}|dur |n|d}|dur,|n|d}ttt |t ||}|S)a Construct an exponential noise schedule. Args: in_sigmas (`torch.Tensor`): The input sigma values to be converted. num_inference_steps (`int`): The number of inference steps to generate the noise schedule for. Returns: `torch.Tensor`: The converted sigma values following an exponential schedule. rDNrEr2r) rr4rDrErCr8expr9mathlog)rFrurtrDrErBrrrry6s     z7FlowMatchEulerDiscreteScheduler._convert_to_exponential333333?alphabetac st|jdr |jjndt|jdr|jjnddur n|ddur,n|dtfddfddd tdd |DD}|S) a Construct a beta noise schedule as proposed in [Beta Sampling is All You Need](https://huggingface.co/papers/2407.12173). Args: in_sigmas (`torch.Tensor`): The input sigma values to be converted. num_inference_steps (`int`): The number of inference steps to generate the noise schedule for. alpha (`float`, *optional*, defaults to `0.6`): The alpha parameter for the beta distribution. beta (`float`, *optional*, defaults to `0.6`): The beta parameter for the beta distribution. Returns: `torch.Tensor`: The converted sigma values following a beta distribution schedule. rDNrEr2rcsg|] }|qSrr)rZppf)rErDrrr^~szDFlowMatchEulerDiscreteScheduler._convert_to_beta..csg|] }tjj|qSr)scipystatsrr)rZrT)rrrrr^sr )rr4rDrErCr8rvr9)rFrurtrrrBr)rrrErDrrzXs       z0FlowMatchEulerDiscreteScheduler._convert_to_betacCs$t|t|d|d|SNr )rrrnrrrrls$z7FlowMatchEulerDiscreteScheduler._time_shift_exponentialcCs||d|d|Srrrnrrrrmsz2FlowMatchEulerDiscreteScheduler._time_shift_linearcCs|jjSrhrirJrrr__len__sz'FlowMatchEulerDiscreteScheduler.__len__)rrFrrrrFNFFFr F)rrh)NNNNN)rr)+rrrr _compatiblesorderrrfloatboolrrGpropertyr"rLrNrPrRrrrgrjrrorsstrrWlistrrrrYr Generatorrtuplerrxryrzrlrmrrrrrr/s'      6     2   k       f'# 0r)r dataclassesrtypingrrrnumpyr8rconfiguration_utilsrrutilsr r r scheduling_utilsr scipy.statsr get_loggerrloggerrrrrrrs