o Gi]@sddlmZddlZddlZddlZddlmZm Z m Z m Z ddl m Z ddlmZddlmZmZmZmZddlmZd d lmZerQddlmmZd Znd ZeeZ d Z!dZ"eGdddeZ#GdddeZ$dS)) dataclassN)CLIPImageProcessorCLIPTextModelWithProjection CLIPTokenizerCLIPVisionModelWithProjection)PriorTransformer)UnCLIPScheduler) BaseOutputis_torch_xla_availableloggingreplace_example_docstring) randn_tensor)DiffusionPipelineTFav Examples: ```py >>> from diffusers import KandinskyPipeline, KandinskyPriorPipeline >>> import torch >>> pipe_prior = KandinskyPriorPipeline.from_pretrained("kandinsky-community/kandinsky-2-1-prior") >>> pipe_prior.to("cuda") >>> prompt = "red cat, 4k photo" >>> out = pipe_prior(prompt) >>> image_emb = out.image_embeds >>> negative_image_emb = out.negative_image_embeds >>> pipe = KandinskyPipeline.from_pretrained("kandinsky-community/kandinsky-2-1") >>> pipe.to("cuda") >>> image = pipe( ... prompt, ... image_embeds=image_emb, ... negative_image_embeds=negative_image_emb, ... height=768, ... width=768, ... num_inference_steps=100, ... ).images >>> image[0].save("cat.png") ``` a Examples: ```py >>> from diffusers import KandinskyPriorPipeline, KandinskyPipeline >>> from diffusers.utils import load_image >>> import PIL >>> import torch >>> from torchvision import transforms >>> pipe_prior = KandinskyPriorPipeline.from_pretrained( ... "kandinsky-community/kandinsky-2-1-prior", torch_dtype=torch.float16 ... ) >>> pipe_prior.to("cuda") >>> img1 = load_image( ... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" ... "/kandinsky/cat.png" ... ) >>> img2 = load_image( ... "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" ... "/kandinsky/starry_night.jpeg" ... ) >>> images_texts = ["a cat", img1, img2] >>> weights = [0.3, 0.3, 0.4] >>> image_emb, zero_image_emb = pipe_prior.interpolate(images_texts, weights) >>> pipe = KandinskyPipeline.from_pretrained("kandinsky-community/kandinsky-2-1", torch_dtype=torch.float16) >>> pipe.to("cuda") >>> image = pipe( ... "", ... image_embeds=image_emb, ... negative_image_embeds=zero_image_emb, ... height=768, ... width=768, ... num_inference_steps=150, ... ).images[0] >>> image.save("starry_cat.png") ``` c@s2eZdZUdZejejBed<ejejBed<dS)KandinskyPriorPipelineOutputa Output class for KandinskyPriorPipeline. Args: image_embeds (`torch.Tensor`) clip image embeddings for text prompt negative_image_embeds (`list[PIL.Image.Image]` or `np.ndarray`) clip image embeddings for unconditional tokens image_embedsnegative_image_embedsN) __name__ __module__ __qualname____doc__torchTensornpndarray__annotations__rrj/home/ubuntu/.local/lib/python3.10/site-packages/diffusers/pipelines/kandinsky/pipeline_kandinsky_prior.pyrxs  rcs`eZdZdZdgZdZdededede de de f fd d Z e ee   d(deeejjBe jBdeededede jee jBd Bde jd Bded BdedefddZddZd)ddZ d*dd Ze ee  ! "d+d#eeeBdeeeBd Bdedede jee jBd Bde jd Bded$ed Bd%efd&d'ZZ S),KandinskyPriorPipelinea Pipeline for generating image prior for Kandinsky 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: prior ([`PriorTransformer`]): The canonical unCLIP prior to approximate the image embedding from the text embedding. image_encoder ([`CLIPVisionModelWithProjection`]): Frozen image-encoder. text_encoder ([`CLIPTextModelWithProjection`]): Frozen text-encoder. tokenizer (`CLIPTokenizer`): Tokenizer of class [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). scheduler ([`UnCLIPScheduler`]): A scheduler to be used in combination with `prior` to generate image embedding. priorztext_encoder->prior image_encoder text_encoder tokenizer schedulerimage_processorcs$t|j||||||ddS)N)r r"r#r$r!r%)super__init__register_modules)selfr r!r"r#r$r% __class__rrr's  zKandinskyPriorPipeline.__init__N@images_and_promptsweightsnum_images_per_promptnum_inference_steps generatorlatentsnegative_prior_promptnegative_promptguidance_scalec  Cs8| p|j} t|t|krtdt|dt|dg} t||D]S\} } t| tr9|| |||||| dj}n6t| tjjt j frft| tjjr^|j | ddj d dj|jj| d} || d }n td t| | || q#t | jdd d }|||||||| d}|d kr|jn|j}t||dS)a Function invoked when using the prior pipeline for interpolation. Args: images_and_prompts (`list[str | PIL.Image.Image | torch.Tensor]`): list of prompts and images to guide the image generation. weights: (`list[float]`): list of weights for each condition in `images_and_prompts` num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. 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. 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`. negative_prior_prompt (`str`, *optional*): The prompt not to guide the prior diffusion process. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). negative_prompt (`str` or `list[str]`, *optional*): The prompt not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). 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. Examples: Returns: [`KandinskyPriorPipelineOutput`] or `tuple` z`images_and_prompts` contains z items and `weights` contains z, items - they should be lists of same length)r3r2r4r5r7r8pt)return_tensorsr)dtypedevicerzq`images_and_prompts` can only contains elements to be of type `str`, `PIL.Image.Image` or `torch.Tensor` but is T)dimkeepdimr.rr)r<len ValueErrorzip isinstancestrrPILImagerrr% pixel_values unsqueezetor!r;typeappendcatsumrr)r)r0r1r2r3r4r5r6r7r8r<image_embeddingscondweight image_embout_zerozero_image_embrrr interpolatesZ 6      z"KandinskyPriorPipeline.interpolatecCsR|dur t||||d}n|j|krtd|jd|||}||j}|S)N)r4r<r;zUnexpected latents shape, got z , expected )rshaperArIinit_noise_sigma)r)rUr;r<r4r5r$rrrprepare_latents s   z&KandinskyPriorPipeline.prepare_latentscCsR|p|j}tdd|jjj|jjjj||jjd}||d}||d}|S)Nr,r)r<r;r) r<rzerosr!config image_sizerIr;repeat)r) batch_sizer<zero_imgrSrrrget_zero_embed+s  z%KandinskyPriorPipeline.get_zero_embedc Cszt|tr t|nd}|j|d|jjddd}|j}|j|} |j|dddj} | j d|j dkrft || sf|j | dd|jjddf} t d |jjd | |ddd|jjf}|||} | j} | j}| j|d d } |j|d d }| j|d d } |r8|durd g|}n;t|t|urtdt|dt|dt|tr|g}n|t|krtd|dt|d|d|d |}|j|d|jjddd}|j|}||j|}|j}|j}|j d}|d|}||||}|j d}|d|d}||||d}|j|d d }t || g} t ||g}t || g} | || fS)Nr, max_lengthTr9)paddingr_ truncationr:longest)r`r:z\The following part of your input was truncated because CLIP can only handle sequences up to z tokens: r)r=r.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`.)rClistr@r#model_max_length input_idsattention_maskboolrIrUrequal batch_decodeloggerwarningr" text_embedslast_hidden_staterepeat_interleaverJ TypeErrorrDrAr[viewrL)r)promptr<r2do_classifier_free_guidancer7r\ text_inputstext_input_ids text_maskuntruncated_ids removed_texttext_encoder_output prompt_embedstext_encoder_hidden_states uncond_tokens uncond_inputuncond_text_mask*negative_prompt_embeds_text_encoder_outputnegative_prompt_embeds!uncond_text_encoder_hidden_statesseq_lenrrr_encode_prompt4s $         z%KandinskyPriorPipeline._encode_promptr9Trs output_type return_dictc Cs\t|tr |g}nt|tstdt|t|tr |g}nt|ts2|dur2tdt||dur>||}d|}|j} t|} | |} |dk} ||| || |\} }}|jj || d|jj }|j j j }|| |f| j| |||j}t||D]P\}}| rt|gdn|}|j ||| ||dj}| r|d\}}||||}|d|jd krd}n||d}|jj|||||d j}trtq|j |}|}|dur|j|jd |jd}|n|d\}}t |d r|j!dur|j"#|d vrtd ||dkr!|$%}|$%}| s(||fSt&||dS)a Function invoked when calling the pipeline for generation. Args: prompt (`str` or `list[str]`): The prompt or prompts to guide the image generation. negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. 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. 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. 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`. 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. output_type (`str`, *optional*, defaults to `"pt"`): The output format of the generate image. Choose between: `"np"` (`np.array`) or `"pt"` (`torch.Tensor`). return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. Examples: Returns: [`KandinskyPriorPipelineOutput`] or `tuple` z2`prompt` has to be of type `str` or `list` but is Nz;`negative_prompt` has to be of type `str` or `list` but is rg?)r<)timestepproj_embeddingencoder_hidden_statesrhr,r)rsampler4 prev_timestepfinal_offload_hook)r9rzBOnly the output types `pt` and `np` are supported not output_type=rr?)'rCrDrerArJ_execution_devicer@rr$ set_timesteps timestepsr rY embedding_dimrWr; enumerate progress_barrrLpredicted_image_embeddingchunkrUstep prev_sample XLA_AVAILABLExm mark_steppost_process_latentsr^r<maybe_free_model_hookshasattrr prior_hookoffloadcpunumpyr)r)rsr7r2r3r4r5r8rrr<r\rtr{r|rwprior_timesteps_tensorritlatent_model_inputr predicted_image_embedding_uncondpredicted_image_embedding_textrrN zero_embedsrrr__call__s 5               zKandinskyPriorPipeline.__call__)r,r-NNNr.r/N)r,N)N)Nr,r-NNr/r9T)!rrrr_exclude_from_cpu_offloadmodel_cpu_offload_seqrrrrr rr'rno_gradr EXAMPLE_INTERPOLATE_DOC_STRINGrerDrErFrfloatint GeneratorrTrWr^rEXAMPLE_DOC_STRINGrir __classcell__rrr*rrs  j  a   r)% dataclassesrrr PIL.ImagerEr transformersrrrrmodelsr schedulersr utilsr r r r utils.torch_utilsrpipeline_utilsrtorch_xla.core.xla_modelcore xla_modelrr get_loggerrrlrrrrrrrrs(      .