o ۷i+l@sddlmZddlZddlmZmZddlmZddlm Z m Z ddl m Z ddl mZmZmZmZdd lmZd d lmZmZerOddlmmZd Znd ZeeZdZdddZ GdddeeZ!dS))CallableN)T5EncoderModel T5Tokenizer)StableDiffusionLoraLoaderMixin)Kandinsky3UNetVQModel) DDPMScheduler) deprecateis_torch_xla_availableloggingreplace_example_docstring) randn_tensor)DiffusionPipelineImagePipelineOutputTFa Examples: ```py >>> from diffusers import AutoPipelineForText2Image >>> import torch >>> pipe = AutoPipelineForText2Image.from_pretrained( ... "kandinsky-community/kandinsky-3", variant="fp16", torch_dtype=torch.float16 ... ) >>> pipe.enable_model_cpu_offload() >>> prompt = "A photograph of the inside of a subway train. There are raccoons sitting on the seats. One of them is reading a newspaper. The window shows the city in the background." >>> generator = torch.Generator(device="cpu").manual_seed(0) >>> image = pipe(prompt, num_inference_steps=25, generator=generator).images[0] ``` cCsX||d}||ddkr|d7}||d}||ddkr$|d7}||||fS)Nrr)heightwidth scale_factor new_height new_widthrrh/home/ubuntu/vllm_env/lib/python3.10/site-packages/diffusers/pipelines/kandinsky3/pipeline_kandinsky3.pydownscale_height_and_width1s  rc$seZdZdZgdZdedededede f fdd Z d d Z e       d3de jdBde jdBde jdBde jdBfddZddZ      d4ddZeddZeddZeddZe eedd d!dd d"d"dddddd#d ddd$gfd%eeeBd&ed'ed(eeeBdBd)edBd*edBd+edBd,e jee jBdBde jdBde jdBde jdBde jdBd-edBd.ed/eeegdfdBd0eef d1d2ZZ S)5Kandinsky3Pipelineztext_encoder->unet->movq)latents prompt_embedsnegative_prompt_embedsnegative_attention_maskattention_mask tokenizer text_encoderunet schedulermovqcs"t|j|||||ddS)N)r"r#r$r%r&)super__init__register_modules)selfr"r#r$r%r& __class__rrr(Es   zKandinsky3Pipeline.__init__cCs`|r,t||dk||dk<|dd}|ddd|f}|ddd|f}||fS)Nrr)torch zeros_likesummax)r* embeddingsr! cut_contextmax_seq_lengthrrrprocess_embedsSs z!Kandinsky3Pipeline.process_embedsTrNFrrr!r c  Cs|dur|durt|t|urtdt|dt|d|dur&|j}|dur2t|tr2d} n|dur@t|tr@t|} n|jd} d} |dur||j|d| d d d } | j |}| j |} |j || d }|d}| || |\}} || d }|j dur|j j}nd}|j ||d}|j\}}}|d|d}||||d}| |d} |r4|dur4|durdg| }n$t|tr|g}n| t|krtd|dt|d|d| d |}|dur*|j|ddd d d d}|j |}|j |} |j || d }|d}|ddd|jdf}| ddd|jdf} || d }n t|}t| } |ra|jd}|j ||d}|j|jkr`|d|d}|| ||d}| |d} nd}d} ||| | fS)aX Encodes the prompt into text encoder hidden states. Args: prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`, *optional*): torch device to place the resulting embeddings on num_images_per_prompt (`int`, *optional*, defaults to 1): number of images that should be generated per prompt do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): whether to use classifier free guidance or not 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. 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`). 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. 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. attention_mask (`torch.Tensor`, *optional*): Pre-generated attention mask. Must provide if passing `prompt_embeds` directly. negative_attention_mask (`torch.Tensor`, *optional*): Pre-generated negative attention mask. Must provide if passing `negative_prompt_embeds` directly. Nz?`negative_prompt` should be the same type to `prompt`, but got z != .rr max_lengthTpt)paddingr8 truncationreturn_tensors)r!r)dtypedevicer-z`negative_prompt`: z has batch size z, but `prompt`: zT. Please make sure that passed `negative_prompt` matches the batch size of `prompt`.)r:r8r;return_attention_maskr<)type TypeError_execution_device isinstancestrlistlenshaper" input_idstor!r#r5 unsqueezer=repeatview ValueErrorr.r/)r*promptdo_classifier_free_guidancenum_images_per_promptr>negative_promptrr _cut_contextr!r batch_sizer8 text_inputstext_input_idsr=bs_embedseq_len_ uncond_tokens uncond_inputrrr encode_prompt[s*                    z Kandinsky3Pipeline.encode_promptcCsR|dur t||||d}n|j|krtd|jd|||}||j}|S)N) generatorr>r=zUnexpected latents shape, got z , expected )rrHrNrJinit_noise_sigma)r*rHr=r>r]rr%rrrprepare_latentss   z"Kandinsky3Pipeline.prepare_latentsc s |durt|tr |dkrtd|dt|d|dur;tfdd|Ds;tdjdfd d |D|durN|durNtd |d |d |durZ|durZtd|durqt|tsqt|tsqtdt||dur|durtd|d|d |dur|dur|j|jkrtd|jd|jd|dur|durtd|dur|dur|jdd|jkrtd|jddd|jd|dur|durtd|dur|dur|jdd|jkrtd|jddd|jddSdSdS)Nrz5`callback_steps` has to be a positive integer but is z of type r6c3|]}|jvVqdSN_callback_tensor_inputs.0kr*rr   z2Kandinsky3Pipeline.check_inputs..2`callback_on_step_end_tensor_inputs` has to be in , but found cg|] }|jvr|qSrrbrdrgrr z3Kandinsky3Pipeline.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'Cannot forward both `negative_prompt`: z and `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` zLPlease provide `negative_attention_mask` along with `negative_prompt_embeds`rz`negative_prompt_embeds` and `negative_attention_mask` must have the same batch_size and token length when passed directly, but got: `negative_prompt_embeds` z != `negative_attention_mask` z:Please provide `attention_mask` along with `prompt_embeds`z`prompt_embeds` and `attention_mask` must have the same batch_size and token length when passed directly, but got: `prompt_embeds` z != `attention_mask` ) rDintrNrAallrcrErFrH) r*rOcallback_stepsrRrr"callback_on_step_end_tensor_inputsr!r rrgr check_inputssz    zKandinsky3Pipeline.check_inputscC|jSra_guidance_scalergrrrguidance_scaleBz!Kandinsky3Pipeline.guidance_scalecCs |jdkS)NrrurgrrrrPFs z.Kandinsky3Pipeline.do_classifier_free_guidancecCrtra)_num_timestepsrgrrr num_timestepsJrxz Kandinsky3Pipeline.num_timestepsg@ipilrrOnum_inference_stepsrwrRrQrrr] output_type return_dictcallback_on_step_endrrc& s$|dd}|dd}|durtddd|dur tddd|durAtfdd|DsAtd jd fd d |Dd }j}|||| | || | |_|durat|t rad}n|durot|t rot |}n| j d}j |j|||| | || | d \} } } } jrt| | g} t| | g} jj||djj}t||d\}}||d||f| j|||j}tdrЈjdurЈjt ||jj}t |_j|d!}t|D]\}}jrt|gdn|}j||| | ddd}jr| d\}} |d| ||}jj!||||dj"}|durai}!|D] }"t#|"|!|"<q1||||!}#|#d|}|#d| } |#d| } |#d| } |#d| } |t |dks||d|kr|djjdkr|$|dur||dkr|t%jd d}$||$||t&rt'(q| d!vrtd"| | d#ksj)j*|d d$d%}%| d&vr|%d'd'}%|%+dd}%|%,-ddd(d./}%| d)kr0|%}%n|}%1|s|%fWdSt2|%d*WdS1s wYdS)+u6 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. num_inference_steps (`int`, *optional*, defaults to 25): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps` timesteps are used. Must be in descending order. guidance_scale (`float`, *optional*, defaults to 3.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. 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_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. 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. 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. attention_mask (`torch.Tensor`, *optional*): Pre-generated attention mask. Must provide if passing `prompt_embeds` directly. negative_attention_mask (`torch.Tensor`, *optional*): Pre-generated negative attention mask. Must provide if passing `negative_prompt_embeds` directly. 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. callback (`Callable`, *optional*): A function that will be called every `callback_steps` steps during inference. The function will be called with the following arguments: `callback(step: int, timestep: int, latents: torch.Tensor)`. callback_steps (`int`, *optional*, defaults to 1): The frequency at which the `callback` function will be called. If not specified, the callback will be called at every step. 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. cross_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). Examples: Returns: [`~pipelines.ImagePipelineOutput`] or `tuple` callbackNrqz1.0.0zhPassing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`znPassing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`c3r`rarbrdrgrrrhriz.Kandinsky3Pipeline.__call__..rjrkcrlrrbrdrgrrrmrnz/Kandinsky3Pipeline.__call__..Trr)rQr>rRrrrSr!r )r>rtext_encoder_offload_hook)totalrF)encoder_hidden_statesencoder_attention_maskrg?)r]rrrr!r order)r9npr|latentzSOnly the output types `pt`, `pil`, `np` and `latent` are supported not output_type=r)force_not_quantizesample)rr|g?rr|)images)3popr rprNrcrCrsrvrDrErFrGrHr\rPr.catboolr% set_timesteps timestepsrr_r=hasattrroffloadrry progress_bar enumerater$chunkstep prev_samplelocalsupdategetattr XLA_AVAILABLExm mark_stepr&decodeclampcpupermutefloatnumpy numpy_to_pilmaybe_free_model_hooksr)&r*rOr}rwrRrQrrr]rrr!r r~rrrrrkwargsrrqr3r>rTrnum_warmup_stepsritlatent_model_input noise_prednoise_pred_uncondnoise_pred_textcallback_kwargsrfcallback_outputsstep_idximagerrgr__call__Ns `             6        H&zKandinsky3Pipeline.__call__) TrNNNNFNN)NNNNNN)!__name__ __module__ __qualname__model_cpu_offload_seqrcrrrr rr(r5r.no_gradTensorr\r_rspropertyrwrPrzr EXAMPLE_DOC_STRINGrErFror Generatorrrr __classcell__rrr+rr;s   F         r)r)"typingrr. transformersrrloadersrmodelsrr schedulersr utilsr r r r utils.torch_utilsrpipeline_utilsrrtorch_xla.core.xla_modelcore xla_modelrr get_loggerrloggerrrrrrrrs"