o Gi @s<ddlZddlmZmZddlZddlmZmZddlm Z m Z ddl m Z ddl mZddlmZmZdd lmZdd lmZmZmZmZmZmZmZdd lmZd d lmZm Z eriddl!m"m#Z$dZ%ndZ%e&e'Z(dZ)    dde*dBde+ej,BdBde-e*dBde-e.dBfddZ/GdddeeZ0dS)N)AnyCallable) T5TokenizerUMT5EncoderModel)MultiPipelineCallbacksPipelineCallback)VaeImageProcessor)AuraFlowLoraLoaderMixin)AuraFlowTransformer2DModel AutoencoderKL)FlowMatchEulerDiscreteScheduler)USE_PEFT_BACKEND deprecateis_torch_xla_availableloggingreplace_example_docstringscale_lora_layersunscale_lora_layers) randn_tensor)DiffusionPipelineImagePipelineOutputTFa Examples: ```py >>> import torch >>> from diffusers import AuraFlowPipeline >>> pipe = AuraFlowPipeline.from_pretrained("fal/AuraFlow", torch_dtype=torch.float16) >>> pipe = pipe.to("cuda") >>> prompt = "A cat holding a sign that says hello world" >>> image = pipe(prompt).images[0] >>> image.save("aura_flow.png") ``` num_inference_stepsdevice timestepssigmascKs|dur |dur td|dur>dtt|jjv}|s(td|jd|jd ||d||j}t |}||fS|durpdtt|jjv}|sZtd|jd|jd ||d ||j}t |}||fS|j|fd |i||j}||fS) 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 valuesrzThe current scheduler class zx's `set_timesteps` does not support custom timestep schedules. Please check whether you are using the correct scheduler.)rrrzv's `set_timesteps` does not support custom sigmas schedules. Please check whether you are using the correct scheduler.)rrr) ValueErrorsetinspect signature set_timesteps parameterskeys __class__rlen) schedulerrrrrkwargsaccepts_timesteps accept_sigmasrrd/home/ubuntu/.local/lib/python3.10/site-packages/diffusers/pipelines/aura_flow/pipeline_aura_flow.pyretrieve_timesteps@s2  r,c.s(eZdZdZgZdZddgZdedede de d e f fd d Z d:d dZ    d;deeeBdeeeBdededejd Bdejd Bdejd Bdejd Bdejd Bdeded BfddZ dtransformer->vaelatents prompt_embeds tokenizer text_encodervae transformerr'csXt|j|||||dt|ddr dt|jjjdnd|_t |jd|_ dS)N)r0r1r2r3r'r2r)vae_scale_factor) super__init__register_modulesgetattrr&r2configblock_out_channelsr6r image_processor)selfr0r1r2r3r'r%rr+r8s  (zAuraFlowPipeline.__init__Nc s|jddks|jddkr#tdjdd|d|d| durDtfdd| DsDtd jd fd d | D|durW|durWtd |d|d|durc|durctd|durzt|tszt|tsztdt||dur|durtd |d|d|dur|durtd|d|d|dur|durtd|dur|durtd|dur|dur|j|jkrtd|jd|jd|j|jkrtd|jd|jddSdSdS)Nrrz-`height` and `width` have to be divisible by z but are z and .c3s|]}|jvVqdSN_callback_tensor_inputs.0kr>rr+ s  z0AuraFlowPipeline.check_inputs..z2`callback_on_step_end_tensor_inputs` has to be in z , but found csg|] }|jvr|qSrrBrDrGrr+ sz1AuraFlowPipeline.check_inputs..zCannot 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` ) r6rallrC isinstancestrlisttypeshape) r>promptheightwidthnegative_promptr/negative_prompt_embedsprompt_attention_masknegative_prompt_attention_mask"callback_on_step_end_tensor_inputsrrGr+ check_inputssn$   zAuraFlowPipeline.check_inputsTr4rPrSdo_classifier_free_guidancenum_images_per_promptrrTrUrVmax_sequence_length lora_scalec s| durt|tr| |_|jdurtrt|j| dur |j|dur,t|tr,d} n|dur:t|tr:t |} n|j d} | } |dur|j |d| ddd}|d}|j |d dd j }|j d |j d krt ||s|j |dd| dd f}td | d |fdd|D}|jdi|d}|dd |j }||}|jdur|jj}n |jdur|jj}nd}|j|d}|j \}}}|d|d}||||d }||d }||d}|r6|dur6|pd}t|tr|g| n|}|j d} |j |d| ddd}fdd|D}|jdi|d}|dd |j } || }|rb|j d}|j|d}|d|d}|| ||d }| |d } | |d} nd}d} |jdur{t|tr{tr{t|j| |||| fS)a  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`). 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. prompt_attention_mask (`torch.Tensor`, *optional*): Pre-generated attention mask for text embeddings. negative_prompt_embeds (`torch.Tensor`, *optional*): Pre-generated negative text embeddings. negative_prompt_attention_mask (`torch.Tensor`, *optional*): Pre-generated attention mask for negative text embeddings. max_sequence_length (`int`, defaults to 256): Maximum sequence length to use for the prompt. lora_scale (`float`, *optional*): A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded. Nr4rT max_lengthpt) truncationr^paddingreturn_tensors input_idslongest)rarbzZThe following part of your input was truncated because T5 can only handle sequences up to z tokens: ci|] \}}||qSrtorErFvrrr+ 9z2AuraFlowPipeline.encode_prompt..attention_mask)dtypercrfrrgrirkrr+rlZrmr)rKr _lora_scaler1rr_execution_devicerLrMr&rOr0rctorchequal batch_decodeloggerwarningitems unsqueezeexpandror3rhrepeatviewreshaper)r>rPrSrZr[rr/rTrUrVr\r] batch_sizer^ text_inputstext_input_idsuntruncated_ids removed_textrobs_embedseq_len_ uncond_tokens uncond_inputrrkr+ encode_prompts,                 zAuraFlowPipeline.encode_promptc Csz|dur |j||dS||t||jt||jf} t|tr3t||kr3tdt|d|dt| |||d}|S)N)rroz/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.) generatorrro)rhintr6rKrMr&rr) r>r~num_channels_latentsrQrRrorrr.rOrrr+prepare_latentsxs    z AuraFlowPipeline.prepare_latentscCs tddd|jjtjddS)N upcast_vaez1.0.0z`upcast_vae` is deprecated. Please use `pipe.vae.to(torch.float32)`. For more details, please refer to: https://github.com/huggingface/diffusers/pull/12619#issue-3606633695.ro)rr2rhrsfloat32rGrrr+rs zAuraFlowPipeline.upcast_vaecC|jSrA)_guidance_scalerGrrr+guidance_scalezAuraFlowPipeline.guidance_scalecCrrA)_attention_kwargsrGrrr+attention_kwargsrz!AuraFlowPipeline.attention_kwargscCrrA)_num_timestepsrGrrr+ num_timestepsrzAuraFlowPipeline.num_timesteps2g @ipilrrrrQrRr output_type return_dictrcallback_on_step_endrWreturnc* Cs||p |jjj|j}|p|jjj|j}|j||||| | | ||d ||_||_|dur4t|tr4d}n|durBt|t rBt |}n| j d}|j }|j durV|j ddnd}|dk}|j|||||| | | |||d \} } } }|r{tj| | gdd} trd }n|}t|j|||d \}}|jjj}||||||| j|| | } tt |||jjd}t ||_|j|d }t|D]\}}|rt| gd n| } t|d g| j d}!|!j| j | jd}!|j| | |!d|j dd}"|r|"!d \}#}$|#||$|#}"|jj"|"|| ddd} |dur7i}%|D] }&t#|&|%|&<q|||||%}'|'$d| } |'$d| } |t |dksR|d|krV|d|jjdkrV|%tr]t&'qWdn 1siwY|dkrv| }(n9|j(jtj)ko|j(jj*})|)r|+| t,t-|j(j./j} |j(j0| |j(jj1ddd}(|j2j3|(|d}(|4|s|(fSt5|(dS)a{ 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`). height (`int`, *optional*, defaults to self.transformer.config.sample_size * self.vae_scale_factor): The height in pixels of the generated image. This is set to 1024 by default for best results. width (`int`, *optional*, defaults to self.transformer.config.sample_size * self.vae_scale_factor): The width in pixels of the generated image. This is set to 1024 by default for best results. num_inference_steps (`int`, *optional*, defaults to 50): 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 used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. guidance_scale (`float`, *optional*, defaults to 5.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. 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.FloatTensor`, *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.FloatTensor`, *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.FloatTensor`, *optional*): Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. 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_xl.StableDiffusionXLPipelineOutput`] instead of a plain tuple. attention_kwargs (`dict`, *optional*): A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). 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. max_sequence_length (`int` defaults to 256): Maximum sequence length to use with the `prompt`. 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. )rWNr4rscaleg?) rPrSrZr[rr/rTrUrVr\r])dimcpu)r)totalrirF)encoder_hidden_statestimesteprr)rr.r/latent)r)images)6r3r; sample_sizer6rXrrrKrLrMr&rOrrrgetrrscat XLA_AVAILABLEr,r' in_channelsrromaxorderr progress_bar enumeratetensorrzrhrchunksteplocalspopupdatexm mark_stepr2float16 force_upcastrnextiterpost_quant_convr#decodescaling_factorr= postprocessmaybe_free_model_hooksr)*r>rPrSrrrr[rQrRrr.r/rUrTrVr\rrrrrWr~rr]rZtimestep_devicerlatent_channelsnum_warmup_stepsritlatent_model_inputr noise_prednoise_pred_uncondnoise_pred_textcallback_kwargsrFcallback_outputsimageneeds_upcastingrrr+__call__sd          6 + zAuraFlowPipeline.__call__)NNNNN) NTr4NNNNNrYNrA)*__name__ __module__ __qualname____doc___optional_componentsmodel_cpu_offload_seqrCrrr r r r8rXrLrMboolrrsrTensorfloatrrrpropertyrrrno_gradrEXAMPLE_DOC_STRING Generatordictrrrrrtupler __classcell__rrr?r+r-{s F                  r-)NNNN)1r typingrrrs transformersrr callbacksrrr=r loadersr modelsr r schedulersr utilsrrrrrrrutils.torch_utilsrpipeline_utilsrrtorch_xla.core.xla_modelcore xla_modelrr get_loggerrrvrrrLrrMrr,r-rrrr+s@    $      ;