o Gig@s.ddlZddlmZmZddlZddlmZmZmZddl m Z ddl m Z ddl mZmZddlmZmZmZdd lmZdd lmZdd lmZdd lmZmZmZmZmZm Z m!Z!dd l"m#Z#ddl$m%Z%m&Z&m'Z'ddl(m)Z)erddl*m+m,Z-dZ.ndZ.e/e0Z1dZ2Gddde%e'eeZ3dS)N)AnyCallable) CLIPTextModelCLIPTextModelWithProjection CLIPTokenizer)CLIPTextModelOutput)VaeImageProcessor)StableDiffusionLoraLoaderMixinTextualInversionLoaderMixin) AutoencoderKLPriorTransformerUNet2DConditionModel)get_timestep_embedding)adjust_lora_scale_text_encoder)KarrasDiffusionSchedulers)USE_PEFT_BACKEND deprecateis_torch_xla_availableloggingreplace_example_docstringscale_lora_layersunscale_lora_layers) randn_tensor)DiffusionPipelineImagePipelineOutputStableDiffusionMixin)StableUnCLIPImageNormalizerTFa Examples: ```py >>> import torch >>> from diffusers import StableUnCLIPPipeline >>> pipe = StableUnCLIPPipeline.from_pretrained( ... "fusing/stable-unclip-2-1-l", torch_dtype=torch.float16 ... ) # TODO update model path >>> pipe = pipe.to("cuda") >>> prompt = "a photo of an astronaut riding a horse on mars" >>> images = pipe(prompt).images >>> images[0].save("astronaut_horse.png") ``` c0seZdZUdZddgZdZeed<eed<e ed<e ed<e ed<e ed<eed <e ed <e ed <e ed <eed <dedede de de de d ed e d e d e d effdd Z  dIdeeBdBdejdBfddZ    dJdejdBdejdBdedBfddZ     dKdejdBdejdBdedBdedBfddZddZdd Zd!d"Z   dLd#d$Zd%d&Z  dId'ejd(ed)ejdBd*ejdBfd+d,Ze e!e"    - .  / 0     1 2  /  3 4 5  dMd6e#e$e#BdBd7edBd8edBd9ed:ed;e#e$e#BdBdejdBdejdBdejdBd?e#dBd@e%dAe&eeejgdfdBdBedCe'e#e(fdBd(edDedEedFejdBdedBf,dGdHZ)Z*S)NStableUnCLIPPipelineac Pipeline for text-to-image generation using stable unCLIP. This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods implemented for all pipelines (downloading, saving, running on a particular device, etc.). The pipeline also inherits the following loading methods: - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings - [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for loading LoRA weights - [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for saving LoRA weights Args: prior_tokenizer ([`CLIPTokenizer`]): A [`CLIPTokenizer`]. prior_text_encoder ([`CLIPTextModelWithProjection`]): Frozen [`CLIPTextModelWithProjection`] text-encoder. prior ([`PriorTransformer`]): The canonical unCLIP prior to approximate the image embedding from the text embedding. prior_scheduler ([`KarrasDiffusionSchedulers`]): Scheduler used in the prior denoising process. image_normalizer ([`StableUnCLIPImageNormalizer`]): Used to normalize the predicted image embeddings before the noise is applied and un-normalize the image embeddings after the noise has been applied. image_noising_scheduler ([`KarrasDiffusionSchedulers`]): Noise schedule for adding noise to the predicted image embeddings. The amount of noise to add is determined by the `noise_level`. tokenizer ([`CLIPTokenizer`]): A [`CLIPTokenizer`]. text_encoder ([`CLIPTextModel`]): Frozen [`CLIPTextModel`] text-encoder. unet ([`UNet2DConditionModel`]): A [`UNet2DConditionModel`] to denoise the encoded image latents. scheduler ([`KarrasDiffusionSchedulers`]): A scheduler to be used in combination with `unet` to denoise the encoded image latents. vae ([`AutoencoderKL`]): Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. priorimage_normalizerz+text_encoder->prior_text_encoder->unet->vaeprior_tokenizerprior_text_encoderprior_schedulerimage_noising_scheduler tokenizer text_encoderunet schedulervaec sdt|j||||||||| | | d t|ddr&dt|jjjdnd|_t |jd|_ dS)N) r#r$r!r%r"r&r'r(r)r*r+r+rr)vae_scale_factor) super__init__register_modulesgetattrlenr+configblock_out_channelsr-r image_processor) selfr#r$r!r%r"r&r'r(r)r*r+ __class__o/home/ubuntu/.local/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip.pyr/s (zStableUnCLIPPipeline.__init__Ntext_model_outputtext_attention_maskcCs(|duryt|tr t|nd}|j|d|jjddd}|j} |j|} |j|dddj} | j d| j dkrjt | | sj|j | dd|jjddf} t d |jjd | | ddd|jjf} || |} | j}| j}n|d j d }|d |d}}|} |j|d d }|j|d d }| j|d d } |rd g|}|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_lengthTptpaddingr= truncationreturn_tensorslongestr@rB\The following part of your input was truncated because CLIP can only handle sequences up to tokens: r)dim) isinstancelistr2r#model_max_length input_idsattention_maskbooltoshapetorchequal batch_decodeloggerwarningr$ text_embedslast_hidden_staterepeat_interleaverepeatviewcat)r6promptdevicenum_images_per_promptdo_classifier_free_guidancer;r< batch_size text_inputstext_input_ids text_maskuntruncated_ids removed_textprior_text_encoder_output prompt_embedstext_enc_hid_states uncond_tokens uncond_inputuncond_text_mask0negative_prompt_embeds_prior_text_encoder_outputnegative_prompt_embedsuncond_text_enc_hid_statesseq_lenr9r9r:_encode_prior_prompts~        z)StableUnCLIPPipeline._encode_prior_promptrhrn lora_scalec KsPd} tdd| dd|jd ||||||||d| } t| d| dg}|S) Nz`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple.z_encode_prompt()1.0.0F standard_warn)r]r^r_r`negative_promptrhrnrrrrr9)r encode_promptrRr\) r6r]r^r_r`rvrhrnrrkwargsdeprecation_messageprompt_embeds_tupler9r9r:_encode_prompts    z#StableUnCLIPPipeline._encode_prompt clip_skipc  Cs|durt|tr||_tst|j|nt|j||dur't|tr'd} n|dur5t|tr5t |} n|j d} |durt|t rJ| ||j }|j |d|j jddd} | j} |j |ddd j} | j d | j d krt| | s|j | dd|j jdd f}td |j jd |t|jjd r|jjjr| j|}nd}| dur|j| ||d}|d}n|j| ||dd}|d | d }|jj|}|jdur|jj}n |jdur|jj}n|j}|j||d}|j \}}}|d|d}||||d }|r|dur|durdg| }nC|dur8t |t |ur8t!dt |dt |dt|trB|g}n| t |kr\t"d|dt |d|d| d |}t|t rk| ||j }|j d}|j |d|ddd}t|jjd r|jjjr|j|}nd}|j|j||d}|d}|r|j d}|j||d}|d|d}|| ||d }|jdurt|trtrt#|j|||fS)a Encodes the prompt into text encoder hidden states. Args: prompt (`str` or `list[str]`, *optional*): prompt to be encoded device: (`torch.device`): torch device num_images_per_prompt (`int`): number of images that should be generated per prompt do_classifier_free_guidance (`bool`): 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. 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. lora_scale (`float`, *optional*): A LoRA scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded. clip_skip (`int`, *optional*): Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that the output of the pre-final layer will be used for computing the prompt embeddings. Nrrr=Tr>r?rCrDrErFrGuse_attention_mask)rN)rNoutput_hidden_states)dtyper^rI?`negative_prompt` should be the same type to `prompt`, but got  != .z`negative_prompt`: z has batch size z, but `prompt`: zT. Please make sure that passed `negative_prompt` matches the batch size of `prompt`.)$rJr _lora_scalerrr(rstrrKr2rQr maybe_convert_promptr'rLrMrRrSrTrUrVhasattrr3r}rNrP text_modelfinal_layer_normrr)rZr[type TypeError ValueErrorr)r6r]r^r_r`rvrhrnrrr|rarbrcrerfrNprompt_embeds_dtypebs_embedrp_rjr=rkr9r9r:rw)s+                     z"StableUnCLIPPipeline.encode_promptcCsnd}tdd|ddd|jjj|}|jj|ddd}|d d dd}|dd d d }|S) Nz{The decode_latents method is deprecated and will be removed in 1.0.0. Please use VaeImageProcessor.postprocess(...) insteaddecode_latentsrsFrtr return_dictrrg?r) rr+r3scaling_factordecodeclampcpupermutefloatnumpy)r6latentsryimager9r9r:rsz#StableUnCLIPPipeline.decode_latentscCXdtt|jjjv}i}|r||d<dtt|jjjv}|r*||d<|SNeta generator)setinspect signaturer%step parameterskeysr6rr accepts_etaextra_step_kwargsaccepts_generatorr9r9r:prepare_prior_extra_step_kwargsz4StableUnCLIPPipeline.prepare_prior_extra_step_kwargscCrr)rrrr*rrrrr9r9r:prepare_extra_step_kwargsrz.StableUnCLIPPipeline.prepare_extra_step_kwargsc Cs|ddks |ddkrtd|d|d|dus(|dur5t|tr(|dkr5td|dt|d|durA|durAtd|durM|durMtd |durdt|tsdt|tsdtd t||durp|durptd |dur|durt|t|urtd t|d t|d|dur|dur|j|jkrtd|jd|jd|dks||jj j krtd|jj j dddS)Nr,rz7`height` and `width` have to be divisible by 8 but are z and rz5`callback_steps` has to be a positive integer but is z of type z[Provide either `prompt` or `prompt_embeds`. Please make sure to define only 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 zProvide either `negative_prompt` or `negative_prompt_embeds`. Cannot leave both `negative_prompt` and `negative_prompt_embeds` undefined.rrzu`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but got: `prompt_embeds` z != `negative_prompt_embeds` z$`noise_level` must be between 0 and rz , inclusive.) rrJintrrrKrrQr&r3num_train_timesteps) r6r]heightwidthcallback_steps noise_levelrvrhrnr9r9r: check_inputssX   z!StableUnCLIPPipeline.check_inputscCsR|dur t||||d}n|j|krtd|jd|||}||j}|S)Nrr^rzUnexpected latents shape, got z , expected )rrQrrPinit_noise_sigma)r6rQrr^rrr*r9r9r:prepare_latentsLs   z$StableUnCLIPPipeline.prepare_latents image_embedsrnoisercCs|durt|j||j|jd}tj|g|jd|jd}|j|j|j|}|j j |||d}|j |}t ||jdddd}||j}t ||fd }|S) aG Add noise to the image embeddings. The amount of noise is controlled by a `noise_level` input. A higher `noise_level` increases the variance in the final un-noised images. The noise is applied in two ways: 1. A noise schedule is applied directly to the embeddings. 2. A vector of sinusoidal time embeddings are appended to the output. In both cases, the amount of noise is controlled by the same `noise_level`. The embeddings are normalized before the noise is applied and un-normalized after the noise is applied. Nrrr^) timestepsrrET)r embedding_dimflip_sin_to_cosdownscale_freq_shiftr)rrQr^rrRtensorr"rPscaler& add_noiseunscalerr\)r6rrrrr9r9r:noise_image_embeddingsWs   z+StableUnCLIPPipeline.noise_image_embeddings$@rpilTr@r]rrnum_inference_stepsguidance_scalervr_rr output_typercallbackrcross_attention_kwargsprior_num_inference_stepsprior_guidance_scale prior_latentsc2 Cs|p |jjj|j}|p|jjj|j}|j||||||| | d|dur-t|tr-d}n|dur;t|tr;t|}n| j d}||}|j }|dk}|j ||||d\}}}|j j ||d|j j}|jjj}|||f|j|| ||j }|| |}t||D]V\} }!|rt|gdn|}"|j |"|!}"|j|"|!|||d j}#|r|#d\}$}%|$||%|$}#|j j|#f|!|d |d d id}|dur| |dkr|| |!|q|j|}|}&|dk}'|dur|d dnd}(|j||||'|| | |(|d \} } |'r t| | g} |j|&|| d}&|'r#t |&} t| |&g}&|j!j ||d|j!j})|jjj"}*||*t#||jt#||jf}+|j|+| j|| | |j!d} |$| |},t||)D]l\} }!|'rmt| gdn| }"|j!|"|!}"|j|"|!| |&|d dd}-|'r|-d\}.}/|.||/|.}-|j!j|-|!| fi|,d d id} |dur| |dkr| t%|j!dd}0||0|!| t&rt'(q^| dks|j)j*| |j)jj+d dd}1n| }1|j,j-|1| d}1|.|s|1fSt/|1dS)u+ The call function to the pipeline for generation. Args: prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The height in pixels of the generated image. width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): The width in pixels of the generated image. num_inference_steps (`int`, *optional*, defaults to 20): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. guidance_scale (`float`, *optional*, defaults to 10.0): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide what to not include in image generation. If not defined, you need to pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) from the [DDIM](https://huggingface.co/papers/2010.02502) paper. Only applies to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. generator (`torch.Generator` or `list[torch.Generator]`, *optional*): A [`torch.Generator`](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 is 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 (prompt weighting). If not provided, text embeddings are generated from the `prompt` input argument. negative_prompt_embeds (`torch.Tensor`, *optional*): Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. output_type (`str`, *optional*, defaults to `"pil"`): The output format of the generated image. Choose between `PIL.Image` or `np.array`. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple. callback (`Callable`, *optional*): A function that calls every `callback_steps` steps during inference. The function is 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 is called. If not specified, the callback is called at every step. cross_attention_kwargs (`dict`, *optional*): A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). noise_level (`int`, *optional*, defaults to `0`): The amount of noise to add to the image embeddings. A higher `noise_level` increases the variance in the final un-noised images. See [`StableUnCLIPPipeline.noise_image_embeddings`] for more details. prior_num_inference_steps (`int`, *optional*, defaults to 25): The number of denoising steps in the prior denoising process. More denoising steps usually lead to a higher quality image at the expense of slower inference. prior_guidance_scale (`float`, *optional*, defaults to 4.0): A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. prior_latents (`torch.Tensor`, *optional*): Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image embedding generation in the prior denoising process. Can be used to tweak the same generation with different prompts. If not provided, a latents tensor is generated by sampling using the supplied random `generator`. clip_skip (`int`, *optional*): Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that the output of the pre-final layer will be used for computing the prompt embeddings. Examples: Returns: [`~pipelines.ImagePipelineOutput`] or `tuple`: [`~ pipeline_utils.ImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated images. )r]rrrrrvrhrnNrrg?)r]r^r_r`rr)timestepproj_embeddingencoder_hidden_statesrN)rsamplerFr) r]r^r_r`rvrhrnrrr|)rrr)rQrr^rrr*)r class_labelsrrorderlatentr)r)images)0r)r3 sample_sizer-rrJrrKr2rQ_execution_devicerqr% set_timestepsrr!rrrr enumerate progress_barrRr\scale_model_inputpredicted_image_embeddingchunkrpost_process_latentsgetrwr zeros_liker* in_channelsrrr1 XLA_AVAILABLExm mark_stepr+rrr5 postprocessmaybe_free_model_hooksr)2r6r]rrrrrvr_rrrrhrnrrrrrrrrrr|rar^!prior_do_classifier_free_guidanceprior_prompt_embeds prior_text_encoder_hidden_statesprior_text_maskprior_timesteps_tensorrprior_extra_step_kwargsitlatent_model_inputr predicted_image_embedding_uncondpredicted_image_embedding_textrr`text_encoder_lora_scalernum_channels_latentsrQr noise_prednoise_pred_uncondnoise_pred_textstep_idxrr9r9r:__call__s h               $    zStableUnCLIPPipeline.__call__)NN)NNNN)NNNNN)NNN)NNNrrNrrNNNNrTNrNrrrNN)+__name__ __module__ __qualname____doc___exclude_from_cpu_offloadmodel_cpu_offload_seqr__annotations__rr rrrrr r/rtuplerRTensorrqrr{rrwrrrrr Generatorrno_gradrEXAMPLE_DOC_STRINGrrKrOrdictrr __classcell__r9r9r7r:r Fs> &    ,  d '  8  = .     r )4rtypingrrrR transformersrrr&transformers.models.clip.modeling_cliprr5r loadersr r modelsr r rmodels.embeddingsr models.lorar schedulersrutilsrrrrrrrutils.torch_utilsrpipeline_utilsrrrstable_unclip_image_normalizerrtorch_xla.core.xla_modelcore xla_modelrr get_loggerrrUrr r9r9r9r:s0     $