o
    Gih                  
   @   s  d dl Z d dlZd dlZd dlZd dlmZ d dlmZ d dl	Z	d dl
mZmZmZ ddlmZmZ ddlmZ ddlmZ ddlmZ dd	lmZ dd
lmZ ddlmZmZmZmZm Z m!Z!m"Z" ddl#m$Z$ ddl%m&Z&m'Z' e  rd dl(m)  m*Z+ dZ,ndZ,e!-e.Z/e rd dl0m1Z1 e rd dl2Z2dZ3				dde4dB de5e	j6B dB de7e4 dB de7e8 dB fddZ9G dd de&Z:G dd de:Z;dS )    N)Callable)GemmaPreTrainedModelGemmaTokenizerGemmaTokenizerFast   )MultiPipelineCallbacksPipelineCallback)VaeImageProcessor)AutoencoderKL)get_2d_rotary_pos_embed_lumina)LuminaNextDiT2DModel)FlowMatchEulerDiscreteScheduler)BACKENDS_MAPPING	deprecateis_bs4_availableis_ftfy_availableis_torch_xla_availableloggingreplace_example_docstring)randn_tensor   )DiffusionPipelineImagePipelineOutputTF)BeautifulSoupaA  
    Examples:
        ```py
        >>> import torch
        >>> from diffusers import LuminaPipeline

        >>> pipe = LuminaPipeline.from_pretrained("Alpha-VLLM/Lumina-Next-SFT-diffusers", torch_dtype=torch.bfloat16)
        >>> # Enable memory optimizations.
        >>> pipe.enable_model_cpu_offload()

        >>> prompt = "Upper body of a young woman in a Victorian-era outfit with brass goggles and leather straps. Background shows an industrial revolution cityscape with smoky skies and tall, metal structures"
        >>> image = pipe(prompt).images[0]
        ```
num_inference_stepsdevice	timestepssigmasc                 K   s  |dur|durt d|dur>dtt| jj v }|s(t d| j d| jd||d| | j}t	|}||fS |durpdtt| jj v }|sZt d| j d| jd||d	| | j}t	|}||fS | j|fd
|i| | j}||fS )a  
    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.

    Args:
        scheduler (`SchedulerMixin`):
            The scheduler to get timesteps from.
        num_inference_steps (`int`):
            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
            must be `None`.
        device (`str` or `torch.device`, *optional*):
            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
        timesteps (`list[int]`, *optional*):
            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
            `num_inference_steps` and `sigmas` must be `None`.
        sigmas (`list[float]`, *optional*):
            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
            `num_inference_steps` and `timesteps` must be `None`.

    Returns:
        `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
        second element is the number of inference steps.
    NzYOnly one of `timesteps` or `sigmas` can be passed. Please choose one to set custom valuesr   zThe current scheduler class zx's `set_timesteps` does not support custom timestep schedules. Please check whether you are using the correct scheduler.)r   r   r   zv's `set_timesteps` does not support custom sigmas schedules. Please check whether you are using the correct scheduler.)r   r   r    )

ValueErrorsetinspect	signatureset_timesteps
parameterskeys	__class__r   len)	schedulerr   r   r   r   kwargsaccepts_timestepsaccept_sigmasr   r   ^/home/ubuntu/.local/lib/python3.10/site-packages/diffusers/pipelines/lumina/pipeline_lumina.pyretrieve_timestepsM   s2   r-   c                2       s  e Zd ZdZedZg ZdZddgZ	de
deded	ed
eeB f
 fddZ				dDdeee B dedejdB dedB dedB f
ddZ									dEdeee B dedeee B dedejdB dejdB dejdB dejdB dejdB defddZdd  Z					dFd!d"ZdGd#d$Zd%d& ZdHd'd(Zed)d* Z ed+d, Z!ed-d. Z"e# e$e%dddd/d0dddddddddd1ddd2d3dddgfdeee B d4edB d5edB d6ed7e&deee B d8ee& dedB d9ej'eej' B dB dejdB dejdB dejdB dejdB dejdB d:edB d;eded<ed=e&dB d>edB d?e(eegdf e)B e*B dB d@ee dAe+e,B f.dBdCZ-  Z.S )ILuminaPipelinea  
    Pipeline for text-to-image generation using Lumina-T2I.

    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)

    Args:
        vae ([`AutoencoderKL`]):
            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
        text_encoder ([`GemmaPreTrainedModel`]):
            Frozen Gemma text-encoder.
        tokenizer (`GemmaTokenizer` or `GemmaTokenizerFast`):
            Gemma tokenizer.
        transformer ([`Transformer2DModel`]):
            A text conditioned `Transformer2DModel` to denoise the encoded image latents.
        scheduler ([`SchedulerMixin`]):
            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
    u5   [#®•©™&@·º½¾¿¡§~\)\(\]\[\}\{\|\\/\*]{1,}ztext_encoder->transformer->vaelatentsprompt_embedstransformerr(   vaetext_encoder	tokenizerc                    sn   t    | j|||||d d| _t| jd| _d| _t| dr+| jd ur+| jj	j
nd| _| j| j | _d S )N)r2   r3   r4   r1   r(      )vae_scale_factor   r1      )super__init__register_modulesr6   r	   image_processormax_sequence_lengthhasattrr1   configsample_sizedefault_sample_sizedefault_image_size)selfr1   r(   r2   r3   r4   r&   r   r,   r:      s$   


zLuminaPipeline.__init__   NFpromptnum_images_per_promptr   clean_caption
max_lengthc                 C   sz  |p| j }t|tr|gn|}t|}| j||d}| j|d| jdddd}|j|}| j|dddj|}	|	j	d |j	d kret
||	se| j|	d d | jd	 df }
td
| j d|
  |j|}| j||dd}|jd }| jd ur| jj}n| jd ur| jj}nd }|j||d}|j	\}}}|d	|d	}||| |d}||d	}||| d}||fS )N)rH   r5   Tpt)pad_to_multiple_ofrI   
truncationpaddingreturn_tensorslongest)rM   rN   rE   z]The following part of your input was truncated because Gemma can only handle sequences up to z	 tokens: attention_maskoutput_hidden_statesdtyper   )_execution_device
isinstancestrr'   _text_preprocessingr4   r=   	input_idstoshapetorchequalbatch_decodeloggerwarningrR   r3   hidden_statesrV   r1   repeatview)rC   rF   rG   r   rH   rI   
batch_sizetext_inputstext_input_idsuntruncated_idsremoved_textprompt_attention_maskr0   rV   _seq_lenr   r   r,   _get_gemma_prompt_embeds   sP   
 "




z'LuminaPipeline._get_gemma_prompt_embedsTdo_classifier_free_guidancenegative_promptnegative_prompt_embedsrk   negative_prompt_attention_maskc              
   K   s  |du r| j }t|tr|gn|}|durt|}n|jd }|du r.| j||||
d\}}|r|du r|dur:|nd}t|trF||g n|}|durct|t|urctdt| dt| dt|trl|g}n|t|krtd| d	t| d
| d	| d	|jd }| j	|d|ddd}|j
|}|j|}	| j||	dd}| jj}|jd }|j\}}}|j||d}|d|d}||| |d}|	|d}	|	|| d}	||||	fS )af  
        Encodes the prompt into text encoder hidden states.

        Args:
            prompt (`str` or `list[str]`, *optional*):
                prompt to be encoded
            negative_prompt (`str` or `list[str]`, *optional*):
                The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`
                instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For
                Lumina-T2I, this should be "".
            do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
                whether to use classifier free guidance or not
            num_images_per_prompt (`int`, *optional*, defaults to 1):
                number of images that should be generated per prompt
            device: (`torch.device`, *optional*):
                torch device to place the resulting embeddings on
            prompt_embeds (`torch.Tensor`, *optional*):
                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
                provided, text embeddings will be generated from `prompt` input argument.
            negative_prompt_embeds (`torch.Tensor`, *optional*):
                Pre-generated negative text embeddings. For Lumina-T2I, it's should be the embeddings of the "" string.
            clean_caption (`bool`, defaults to `False`):
                If `True`, the function will preprocess and clean the provided caption before encoding.
            max_sequence_length (`int`, defaults to 256): Maximum sequence length to use for the prompt.
        Nr   )rF   rG   r   rH    z?`negative_prompt` should be the same type to `prompt`, but got z != .z`negative_prompt`: z has batch size z, but `prompt`: zT. Please make sure that passed `negative_prompt` matches the batch size of `prompt`.rE   rI   TrJ   )rM   rI   rL   rN   rQ   rT   rU   rP   )rW   rX   rY   r'   r]   rn   type	TypeErrorr   r4   r[   r\   rR   r3   rV   rc   rd   re   )rC   rF   ro   rp   rG   r   r0   rq   rk   rr   rH   r)   rf   prompt_max_lengthnegative_text_inputsnegative_text_input_idsnegative_dtyperl   rm   r   r   r,   encode_prompt  st   '





zLuminaPipeline.encode_promptc                 C   sX   dt t| jjj v }i }|r||d< dt t| jjj v }|r*||d< |S )Neta	generator)r    r!   r"   r(   stepr$   r%   )rC   r}   r|   accepts_etaextra_step_kwargsaccepts_generatorr   r   r,   prepare_extra_step_kwargst  s   z(LuminaPipeline.prepare_extra_step_kwargsc
           
         s  | j d  dks| j d  dkr#td j d  d| d| d|	d urDt fdd|	D sDtd	 j d
 fdd|	D  |d urW|d urWtd| d| d|d u rc|d u rctd|d urzt|tszt|tsztdt| |d ur|d urtd| d| d|d ur|d urtd| d| d|d ur|d u rtd|d ur|d u rtd|d ur|d ur|j|jkrtd|j d|j d|j|jkrtd|j d|j dd S d S d S )Nr   r   z-`height` and `width` have to be divisible by z	 but are z and rt   c                 3   s    | ]}| j v V  qd S N_callback_tensor_inputs.0krC   r   r,   	<genexpr>  s    

z.LuminaPipeline.check_inputs.<locals>.<genexpr>z2`callback_on_step_end_tensor_inputs` has to be in z, but found c                    s   g | ]	}| j vr|qS r   r   r   r   r   r,   
<listcomp>  s    z/LuminaPipeline.check_inputs.<locals>.<listcomp>zCannot forward both `prompt`: z and `prompt_embeds`: z2. Please make sure to only forward one of the two.zeProvide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined.z2`prompt` has to be of type `str` or `list` but is z and `negative_prompt_embeds`: z'Cannot forward both `negative_prompt`: zEMust provide `prompt_attention_mask` when specifying `prompt_embeds`.zWMust provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.zu`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but got: `prompt_embeds` z != `negative_prompt_embeds` z`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but got: `prompt_attention_mask` z% != `negative_prompt_attention_mask` )	r6   r   allr   rX   rY   listru   r]   )
rC   rF   heightwidthrp   r0   rq   rk   rr   "callback_on_step_end_tensor_inputsr   r   r,   check_inputs  sn   $zLuminaPipeline.check_inputsc                    s    rt  sttd d d td d  r0t s0ttd d d td d t|ttfs:|g}dt	f fdd	fd
d|D S )Nbs4rP   zSetting `clean_caption=True`z#Setting `clean_caption` to False...Fftfytextc                    s,    r | }  | } | S |   } | S r   )_clean_captionlowerstrip)r   )rH   rC   r   r,   process  s   

z3LuminaPipeline._text_preprocessing.<locals>.processc                    s   g | ]} |qS r   r   )r   t)r   r   r,   r     s    z6LuminaPipeline._text_preprocessing.<locals>.<listcomp>)
r   ra   rb   r   formatr   rX   tupler   rY   )rC   r   rH   r   )rH   r   rC   r,   rZ     s   



z"LuminaPipeline._text_preprocessingc                 C   s  t |}t|}|  }tdd|}tdd|}tdd|}t|ddj}tdd|}td	d|}td
d|}tdd|}tdd|}tdd|}tdd|}tdd|}tdd|}tdd|}tdd|}tdd|}tdd|}tdd|}tdd|}tdd|}tdd|}tdd|}tdd|}tdd|}td d|}td!d|}t| j	d|}td"d|}t
d#}tt||d$krt|d|}t|}tt|}td%d|}td&d|}td'd|}td(d|}td)d|}td*d|}td+d|}td,d|}td-d|}td.d|}td/d0|}td1d2|}td3d|}|  td4d5|}td6d|}td7d|}td8d|}| S )9Nz<person>personzk\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))rs   zh\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))zhtml.parser)featuresz
@[\w\d]+\bz[\u31c0-\u31ef]+z[\u31f0-\u31ff]+z[\u3200-\u32ff]+z[\u3300-\u33ff]+z[\u3400-\u4dbf]+z[\u4dc0-\u4dff]+z[\u4e00-\u9fff]+z|[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+-u   [`´«»“”¨]"u   [‘’]'z&quot;?z&ampz"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3} z\d:\d\d\s+$z\\nz
#\d{1,3}\bz	#\d{5,}\bz
\b\d{6,}\bz0[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)z
[\"\']{2,}z[\.]{2,}z\s+\.\s+z	(?:\-|\_)r   z\b[a-zA-Z]{1,3}\d{3,15}\bz\b[a-zA-Z]+\d+[a-zA-Z]+\bz\b\d+[a-zA-Z]+\d+\bz!(worldwide\s+)?(free\s+)?shippingz(free\s)?download(\sfree)?z\bclick\b\s(?:for|on)\s\w+z9\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?z\bpage\s+\d+\bz*\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\bu   \b\d+\.?\d*[xх×]\d+\.?\d*\bz
\b\s+\:\s+z: z(\D[,\./])\bz\1 z\s+z^[\"\']([\w\W]+)[\"\']$z\1z^[\'\_,\-\:;]z[\'\_,\-\:\-\+]$z^\.\S+$)rY   ulunquote_plusr   r   resubr   r   bad_punct_regexcompiler'   findallr   fix_texthtmlunescape)rC   captionregex2r   r   r,   r     s   
	

zLuminaPipeline._clean_captionc	           
      C   sz   ||t || j t || j f}	t|tr(t||kr(tdt| d| d|d u r6t|	|||d}|S ||}|S )Nz/You have passed a list of generators of length z+, but requested an effective batch size of z@. Make sure the batch size matches the length of the generators.)r}   r   rV   )intr6   rX   r   r'   r   r   r\   )
rC   rf   num_channels_latentsr   r   rV   r   r}   r/   r]   r   r   r,   prepare_latentsU  s    
zLuminaPipeline.prepare_latentsc                 C      | j S r   _guidance_scaler   r   r   r,   guidance_scalei     zLuminaPipeline.guidance_scalec                 C   s
   | j dkS )NrE   r   r   r   r   r,   ro   p  s   
z*LuminaPipeline.do_classifier_free_guidancec                 C   r   r   )_num_timestepsr   r   r   r,   num_timestepst  r   zLuminaPipeline.num_timesteps   g      @pilr7         ?r   r   r   r   r   r}   output_typereturn_dictr=   scaling_watershedproportional_attncallback_on_step_endr   returnc           6      C   sx  |p| j | j }|p| j | j }| j|||||||||d	 || _i }|dur/t|tr/d}n|dur=t|tr=t|}n|jd }|rM| j	d d |d< t
|| | j	d  }| j}|dk}| j|||||||||||d	\}}}}|rtj||gdd
}tj||gdd
}trd}n|}t| j|||d\}}| jjj}| || ||||j||	|
}
t|| _| j|dI}t|D ];\} }!|rt|
gd n|
}"|!}#t|#s	|"jjdk}$|"jjdk}%t|#tr|$s|%rtjntj }&n
|$s|%rtj!ntj"}&tj#|#g|&|"jd}#nt|#jdkr|#d $|"j}#|#%|"jd }#d|#| jjj&  }#|#d |k r6|}'d}(nd}'|}(t'| jj(dd|'|(d})| j|"|#|||)|ddd }*|*j)ddd
d }*|r|*ddddf |*ddddf }+},tj*|+t|+d dd
\}-}.|.||-|.   }/tj|/|/gdd
}+tj|+|,gdd
}*|*j)ddd
\}*}0|
j}1|* }*| jj+|*|!|
ddd }
|
j|1krtj,j-. r|
$|1}
|/  |duri }2|D ]
}3t0 |3 |2|3< q|| | |!|2}4|41d|
}
|41d|}trt23  qW d   n	1 sw   Y  |dks+|
| j4jj5 }
| j4j6|
ddd }5| j7j8|5|d}5n|
}5| 9  |s7|5fS t:|5dS )u  
        Function invoked when calling the pipeline for generation.

        Args:
            prompt (`str` or `list[str]`, *optional*):
                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
                instead.
            negative_prompt (`str` or `list[str]`, *optional*):
                The prompt or prompts not to guide the image generation. If not defined, one has to pass
                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
                less than `1`).
            num_inference_steps (`int`, *optional*, defaults to 30):
                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
                expense of slower inference.
            sigmas (`list[float]`, *optional*):
                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
                will be used.
            guidance_scale (`float`, *optional*, defaults to 4.0):
                Guidance scale as defined in [Classifier-Free Diffusion
                Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2.
                of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
                `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
                the text `prompt`, usually at the expense of lower image quality.
            num_images_per_prompt (`int`, *optional*, defaults to 1):
                The number of images to generate per prompt.
            height (`int`, *optional*, defaults to self.unet.config.sample_size):
                The height in pixels of the generated image.
            width (`int`, *optional*, defaults to self.unet.config.sample_size):
                The width in pixels of the generated image.
            eta (`float`, *optional*, defaults to 0.0):
                Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only
                applies to [`schedulers.DDIMScheduler`], will be ignored for others.
            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
                to make generation deterministic.
            latents (`torch.Tensor`, *optional*):
                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
                tensor will be generated by sampling using the supplied random `generator`.
            prompt_embeds (`torch.Tensor`, *optional*):
                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
                provided, text embeddings will be generated from `prompt` input argument.
            prompt_attention_mask (`torch.Tensor`, *optional*): Pre-generated attention mask for text embeddings.
            negative_prompt_embeds (`torch.Tensor`, *optional*):
                Pre-generated negative text embeddings. For Lumina-T2I this negative prompt should be "". If not
                provided, negative_prompt_embeds will be generated from `negative_prompt` input argument.
            negative_prompt_attention_mask (`torch.Tensor`, *optional*):
                Pre-generated attention mask for negative text embeddings.
            output_type (`str`, *optional*, defaults to `"pil"`):
                The output format of the generate image. Choose between
                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether or not to return a [`~pipelines.stable_diffusion.IFPipelineOutput`] instead of a plain tuple.
            clean_caption (`bool`, *optional*, defaults to `True`):
                Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to
                be installed. If the dependencies are not installed, the embeddings will be created from the raw
                prompt.
            max_sequence_length (`int` defaults to 120):
                Maximum sequence length to use with the `prompt`.
            callback_on_step_end (`Callable`, *optional*):
                A function that calls at the end of each denoising steps during the inference. The function is called
                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
                `callback_on_step_end_tensor_inputs`.
            callback_on_step_end_tensor_inputs (`list`, *optional*):
                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
                `._callback_tensor_inputs` attribute of your pipeline class.

        Examples:

        Returns:
            [`~pipelines.ImagePipelineOutput`] or `tuple`:
                If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is
                returned where the first element is a list with the generated images
        )r0   rq   rk   rr   r   NrE   r      r   base_sequence_lengthr   )	rp   rG   r   r0   rq   rk   rr   rH   r=   )dimcpu)r   )totalmpsnpurU   i  )linear_factor
ntk_factorF)rc   timestepencoder_hidden_statesencoder_maskimage_rotary_embcross_attention_kwargsr   r   )r   r/   r0   latent)r   )images);rA   r6   r   r   rX   rY   r   r'   r]   rB   mathsqrtrW   r{   r^   catXLA_AVAILABLEr-   r(   r1   r?   in_channelsr   rV   r   progress_bar	enumerate	is_tensorr   ru   floatfloat32float64int32int64tensorr\   expandnum_train_timestepsr   head_dimchunksplitr~   backendsr   is_availableupdatelocalspopxm	mark_stepr2   scaling_factordecoder<   postprocessmaybe_free_model_hooksr   )6rC   rF   r   r   r   r   rp   r   rG   r}   r/   r0   rq   rk   rr   r   r   rH   r=   r   r   r   r   r   rf   r   r   ro   timestep_devicer   latent_channelsr   ir   latent_model_inputcurrent_timestepis_mpsis_npurV   r   r   r   
noise_prednoise_pred_epsnoise_pred_restnoise_pred_cond_epsnoise_pred_uncond_epsnoise_pred_halfrl   latents_dtypecallback_kwargsr   callback_outputsimager   r   r,   __call__x  s$  h






	*



d
zLuminaPipeline.__call__)rE   NFN)	TNrE   NNNNNF)NNNNN)Fr   )/__name__
__module____qualname____doc__r   r   r   _optional_componentsmodel_cpu_offload_seqr   r   r   r
   r   r   r   r:   rY   r   r   r^   r   boolrn   Tensorr{   r   r   rZ   r   r   propertyr   ro   r   no_gradr   EXAMPLE_DOC_STRINGr   	Generatorr   r   r   r   r   r   __classcell__r   r   rD   r,   r.      s2   

;

	

n

E
r




	
r.   c                       s6   e Zd ZdededededeeB f
 fddZ	  Z
S )LuminaText2ImgPipeliner1   r(   r2   r3   r4   c                    s*   d}t dd| t j|||||d d S )Nz`LuminaText2ImgPipeline` has been renamed to `LuminaPipeline` and will be removed in a future version. Please use `LuminaPipeline` instead.zAdiffusers.pipelines.lumina.pipeline_lumina.LuminaText2ImgPipelinez0.34)r1   r(   r2   r3   r4   )r   r9   r:   )rC   r1   r(   r2   r3   r4   deprecation_messagerD   r   r,   r:     s   
zLuminaText2ImgPipeline.__init__)r  r  r  r   r   r
   r   r   r   r:   r  r   r   rD   r,   r    s    r  )NNNN)<r   r!   r   r   urllib.parseparser   typingr   r^   transformersr   r   r   	callbacksr   r   r<   r	   modelsr
   models.embeddingsr   $models.transformers.lumina_nextdit2dr   
schedulersr   utilsr   r   r   r   r   r   r   utils.torch_utilsr   pipeline_utilsr   r   torch_xla.core.xla_modelcore	xla_modelr   r   
get_loggerr  ra   r   r   r   r  r   rY   r   r   r   r-   r.   r  r   r   r   r,   <module>   s`   $	



;      *