o ۷i@sNddlZddlmZddlZddlZddlZddlm Z ddl m Z m Z mZmZddlmZmZddlmZddlmZdd lmZdd lmZdd lmZdd lmZmZm Z m!Z!dd l"m#Z#ddl$m%Z%ddl&m'Z'er{ddl(m)m*Z+dZ,ndZ,e -e.Z/erddl0Z0e -e.Z/dZ1ddZ2ddZ3ddZ4Gddde%eZ5dS)N)Callable) functional) CLIPTextModel CLIPTokenizer"Qwen2_5_VLForConditionalGenerationQwen2VLProcessor)MultiPipelineCallbacksPipelineCallback)VaeImageProcessor)KandinskyLoraLoaderMixin) AutoencoderKL)Kandinsky5Transformer3DModel)FlowMatchEulerDiscreteScheduler)is_ftfy_availableis_torch_xla_availableloggingreplace_example_docstring) randn_tensor)DiffusionPipeline)KandinskyImagePipelineOutputTFa< Examples: ```python >>> import torch >>> from diffusers import Kandinsky5T2IPipeline >>> # Available models: >>> # kandinskylab/Kandinsky-5.0-T2I-Lite-sft-Diffusers >>> # kandinskylab/Kandinsky-5.0-T2I-Lite-pretrain-Diffusers >>> model_id = "kandinskylab/Kandinsky-5.0-T2I-Lite-sft-Diffusers" >>> pipe = Kandinsky5T2IPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16) >>> pipe = pipe.to("cuda") >>> prompt = "A cat and a dog baking a cake together in a kitchen." >>> output = pipe( ... prompt=prompt, ... negative_prompt="", ... height=1024, ... width=1024, ... num_inference_steps=50, ... guidance_scale=3.5, ... ).frames[0] ``` cCs(trt|}tt|}|S)z Copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/wan/pipeline_wan.py Clean text using ftfy if available and unescape HTML entities. )rftfyfix_texthtmlunescapestriptextr k/home/ubuntu/vllm_env/lib/python3.10/site-packages/diffusers/pipelines/kandinsky5/pipeline_kandinsky_t2i.py basic_cleanWs r"cCstdd|}|}|S)z Copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/wan/pipeline_wan.py Normalize whitespace in text by replacing multiple spaces with single space. z\s+ )resubrrr r r!whitespace_cleancsr&cCstt|}|S)z Copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/wan/pipeline_wan.py Apply both basic cleaning and whitespace normalization to prompts. )r&r"rr r r! prompt_cleanns r'c+seZdZdZdZgdZdededede de d e d e ffd d Z  dAdeedejd Bdedejd BfddZ dBdeeeBdejd Bdejd BfddZ   dCdeeeBdededejd Bdejd Bf ddZ dDddZ    dEded ed!ed"edejd Bdejd Bd#ejeejBd Bd$ejd Bd%ejfd&d'Zed(d)Zed*d+Zed,d-Zee e!d d ddd.d/dd d d d d d d d d0d1d d$gdfdeeeBd2eeeBd Bd!ed"ed3ed4e"ded Bd#ejeejBd Bd$ejd Bd5ejd Bd6ejd Bd7ejd Bd8ejd Bd9ejd Bd:ejd Bd;ed Bdeedef(d?d@Z'Z(S)FKandinsky5T2IPipelinea Pipeline for text-to-image generation using Kandinsky 5.0. 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.). Args: transformer ([`Kandinsky5Transformer3DModel`]): Conditional Transformer to denoise the encoded image latents. vae ([`AutoencoderKL`]): Variational Auto-Encoder Model [black-forest-labs/FLUX.1-dev (vae)](https://huggingface.co/black-forest-labs/FLUX.1-dev) to encode and decode videos to and from latent representations. text_encoder ([`Qwen2_5_VLForConditionalGeneration`]): Frozen text-encoder [Qwen2.5-VL](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct). tokenizer ([`AutoProcessor`]): Tokenizer for Qwen2.5-VL. text_encoder_2 ([`CLIPTextModel`]): Frozen [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant. tokenizer_2 ([`CLIPTokenizer`]): Tokenizer for CLIP. scheduler ([`FlowMatchEulerDiscreteScheduler`]): A scheduler to be used in combination with `transformer` to denoise the encoded image latents. z.text_encoder->text_encoder_2->transformer->vae)latentsprompt_embeds_qwenprompt_embeds_clipnegative_prompt_embeds_qwennegative_prompt_embeds_clip transformervae text_encoder tokenizertext_encoder_2 tokenizer_2 schedulerc sPt|j|||||||dd|_d|_d|_t|jd|_gd|_dS)N)r.r/r0r1r2r3r4z<|im_start|>system You are a promt engineer. Describe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|> <|im_start|>user {}<|im_end|>))vae_scale_factor))r8))r:r9))r<r;))r>r=) super__init__register_modulesprompt_template prompt_template_encode_start_idxvae_scale_factor_spatialr image_processor resolutions)selfr.r/r0r1r2r3r4 __class__r r!r@s  zKandinsky5T2IPipeline.__init__Npromptdevicemax_sequence_lengthdtypec sn|pj}|p jj}fdd|D}j|}j|dddddd}|jd|krft|D]5\}} ||jd } j| |d d} t| d kre| dt|  ||<t d |d | q0j|dd|dddd |} j| ddddddddjdf} | dddjdf}t j |dd d}tj|dd dj t jd}| ||fS)aI Encode prompt using Qwen2.5-VL text encoder. This method processes the input prompt through the Qwen2.5-VL model to generate text embeddings suitable for image generation. Args: prompt list[str]: Input list of prompts device (torch.device): Device to run encoding on max_sequence_length (int): Maximum sequence length for tokenization dtype (torch.dtype): Data type for embeddings Returns: tuple[torch.Tensor, torch.Tensor]: Text embeddings and cumulative sequence lengths csg|]}j|qSr )rBformat.0prGr r! sz=Kandinsky5T2IPipeline._encode_prompt_qwen..Nptlongest)rimagesvideosreturn_tensorspadding input_idsrrzXThe following part of your input was truncated because `max_sequence_length` is set to z tokens: T)rrWrX max_length truncationrYrZ)r[ return_dictoutput_hidden_states hidden_statesattention_maskr)dim)rr)value)rN)_execution_devicer0rNrCr1shape enumeratedecodelenloggerwarningtotorchcumsumsumFpadint32)rGrKrLrMrN full_textsmax_allowed_lenuntruncated_idsirtokens removed_textinputsembedsrc cu_seqlensr rSr!_encode_prompt_qwensj     z)Kandinsky5T2IPipeline._encode_prompt_qwencCsP|p|j}|p |jj}|j|dddddd|}|jdi|d}||S) a Encode prompt using CLIP text encoder. This method processes the input prompt through the CLIP model to generate pooled embeddings that capture semantic information. Args: prompt (str | list[str]): Input prompt or list of prompts device (torch.device): Device to run encoding on dtype (torch.dtype): Data type for embeddings Returns: torch.Tensor: Pooled text embeddings from CLIP MTr^rU)r^r_add_special_tokensrZrY pooler_outputNr )rfr2rNr3rm)rGrKrLrNrz pooled_embedr r r!_encode_prompt_clips   z)Kandinsky5T2IPipeline._encode_prompt_cliprnum_images_per_promptc Cs|p|j}|p |jj}t|ts|g}t|}dd|D}|j||||d\}}|j|||d} |d|d}| ||d|j d}| d|d} | ||d} | } | |} t t jdg|t jd| dg} || | fS) a Encodes a single prompt (positive or negative) into text encoder hidden states. This method combines embeddings from both Qwen2.5-VL and CLIP text encoders to create comprehensive text representations for image generation. Args: prompt (`str` or `list[str]`): Prompt to be encoded. num_images_per_prompt (`int`, *optional*, defaults to 1): Number of images to generate per prompt. max_sequence_length (`int`, *optional*, defaults to 512): Maximum sequence length for text encoding. Must be less than 1024 device (`torch.device`, *optional*): Torch device. dtype (`torch.dtype`, *optional*): Torch dtype. Returns: tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - Qwen text embeddings of shape (batch_size * num_images_per_prompt, sequence_length, embedding_dim) - CLIP pooled embeddings of shape (batch_size * num_images_per_prompt, clip_embedding_dim) - Cumulative sequence lengths (`cu_seqlens`) for Qwen embeddings of shape (batch_size * num_images_per_prompt + 1,) cSsg|]}t|qSr )r'rPr r r!rTJsz7Kandinsky5T2IPipeline.encode_prompt..)rKrLrMrN)rKrLrNrr\rrLrN)rfr0rN isinstancelistrjr}rrepeatviewrgdiffrepeat_interleaverncattensorrsro) rGrKrrMrLrN batch_sizer*prompt_cu_seqlensr+original_lengthsrepeated_lengthsrepeated_cu_seqlensr r r! encode_prompt!sF !      z#Kandinsky5T2IPipeline.encode_promptc  s| dur | dkr td||fjvr-dddjD} td| d|d |d | durNtfd d | DsNtd jdfdd| D|dusZ|dusZ| durj|dusf|dusf| durjtd|dusv|dusv| dur|dus|dus| durtd|dur|durtd|durt|tst|t stdt ||durt|tst|t stdt |dSdSdS)al Validate input parameters for the pipeline. Args: prompt: Input prompt negative_prompt: Negative prompt for guidance height: Image height width: Image width prompt_embeds_qwen: Pre-computed Qwen prompt embeddings prompt_embeds_clip: Pre-computed CLIP prompt embeddings negative_prompt_embeds_qwen: Pre-computed Qwen negative prompt embeddings negative_prompt_embeds_clip: Pre-computed CLIP negative prompt embeddings prompt_cu_seqlens: Pre-computed cumulative sequence lengths for Qwen positive prompt negative_prompt_cu_seqlens: Pre-computed cumulative sequence lengths for Qwen negative prompt callback_on_step_end_tensor_inputs: Callback tensor inputs Raises: ValueError: If inputs are invalid Nr8z*max_sequence_length must be less than 1024,cSs"g|] \}}d|d|dqS)(r)r )rQwhr r r!rTs"z6Kandinsky5T2IPipeline.check_inputs..z'`height` and `width` have to be one of z , but are z and z(. Dimensions will be resized accordinglyc3s|]}|jvVqdS)N_callback_tensor_inputsrQkrSr r! s  z5Kandinsky5T2IPipeline.check_inputs..z2`callback_on_step_end_tensor_inputs` has to be in z , but found csg|] }|jvr|qSr rrrSr r!rTszuIf any of `prompt_embeds_qwen`, `prompt_embeds_clip`, or `prompt_cu_seqlens` is provided, all three must be provided.zIf any of `negative_prompt_embeds_qwen`, `negative_prompt_embeds_clip`, or `negative_prompt_cu_seqlens` is provided, all three must be provided.zProvide either `prompt` or `prompt_embeds_qwen` (and corresponding `prompt_embeds_clip` and `prompt_cu_seqlens`). Cannot leave all undefined.z2`prompt` has to be of type `str` or `list` but is z;`negative_prompt` has to be of type `str` or `list` but is ) ValueErrorrFjoinrkrlallrrstrrtype)rGrKnegative_promptheightwidthr*r+r,r-rnegative_prompt_cu_seqlens"callback_on_step_end_tensor_inputsrMresolutions_strr rSr! check_inputs|sP# z"Kandinsky5T2IPipeline.check_inputsr8rnum_channels_latentsrr generatorr)returnc Cs||dur |j||dS|dt||jt||j|f} t|tr4t||kr4tdt|d|dt| |||d}|S)a Prepare initial latent variables for text-to-image generation. This method creates random noise latents Args: batch_size (int): Number of images to generate num_channels_latents (int): Number of channels in latent space height (int): Height of generated image width (int): Width of generated image dtype (torch.dtype): Data type for latents device (torch.device): Device to create latents on generator (torch.Generator): Random number generator latents (torch.Tensor): Pre-existing latents to use Returns: torch.Tensor: Prepared latent tensor Nrrz/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.)rrLrN)rmintrDrrrjrr) rGrrrrrNrLrr)rgr r r!prepare_latentss    z%Kandinsky5T2IPipeline.prepare_latentscC|jS)z%Get the current guidance scale value.)_guidance_scalerSr r r!guidance_scalez$Kandinsky5T2IPipeline.guidance_scalecCr)z&Get the number of denoising timesteps.)_num_timestepsrSr r r! num_timesteps rz#Kandinsky5T2IPipeline.num_timestepscCr)z)Check if generation has been interrupted.) _interruptrSr r r! interruptrzKandinsky5T2IPipeline.interrupt2g @pilTrnum_inference_stepsrr*r+r,r-rr output_typer`callback_on_step_endrc*st|ttfr |j}|j||| | | | ||||d f|jvr5|jtfdd|jD\||_d|_ |j }|j j }|durQt|t rQd}|g}n|dur_t|tr_t|}n| jd}| duru|j|||||d\} } }|jd kr|durd }t|t r|dur|gt|n|g}nt|t|krtd t|d t|d | dur|j|||||d\} } }|jj||d|jj}|j jj}|j||||||| d} tjd|dtj|jd|dtj|jd|dg}tj||d}|durtj||dnd}gd}d}t|||jj }t||_!|j"|d} t#|D]\}!}"|j$rEq;|"%d&||}#|j | '|| '|| '||#'|||||dd j(}$|jd kr| dur|j | '|| '|| '||#'|||||dd j(}%|%||$|%}$|jj)|$ddddf|"| ddd} |duri}&|D] }'t*|'|&|'<q|||!|"|&}(|(+d| } |(+d| } |(+d| } |(+d| } |(+d| } |!t|dks|!d|kr|!d|jj dkr| ,t-r t./q;Wdn 1swY| ddddddddd|f} |dkr}| '|j0j } | 1||d|j|j|} | 2dddddd} | 1||||j|j} | |j0jj3} |j04| j(})|j5j6|)|d})n| })|7|s|)fSt8|)d S)!aX The call function to the pipeline for text-to-image generation. Args: prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, pass `prompt_embeds` instead. negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts to avoid during image generation. If not defined, pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale` < `1`). height (`int`, defaults to `1024`): The height in pixels of the generated image. width (`int`, defaults to `1024`): The width in pixels of the generated image. num_inference_steps (`int`, defaults to `50`): The number of denoising steps. guidance_scale (`float`, defaults to `5.0`): Guidance scale as defined in classifier-free guidance. 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*): A torch generator to make generation deterministic. latents (`torch.Tensor`, *optional*): Pre-generated noisy latents. prompt_embeds_qwen (`torch.Tensor`, *optional*): Pre-generated Qwen text embeddings. prompt_embeds_clip (`torch.Tensor`, *optional*): Pre-generated CLIP text embeddings. negative_prompt_embeds_qwen (`torch.Tensor`, *optional*): Pre-generated Qwen negative text embeddings. negative_prompt_embeds_clip (`torch.Tensor`, *optional*): Pre-generated CLIP negative text embeddings. prompt_cu_seqlens (`torch.Tensor`, *optional*): Pre-generated cumulative sequence lengths for Qwen positive prompt. negative_prompt_cu_seqlens (`torch.Tensor`, *optional*): Pre-generated cumulative sequence lengths for Qwen negative prompt. output_type (`str`, *optional*, defaults to `"pil"`): The output format of the generated image. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`KandinskyImagePipelineOutput`]. callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*): A function that is called at the end of each denoising step. callback_on_step_end_tensor_inputs (`List`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. max_sequence_length (`int`, defaults to `512`): The maximum sequence length for text encoding. Examples: Returns: [`~KandinskyImagePipelineOutput`] or `tuple`: If `return_dict` is `True`, [`KandinskyImagePipelineOutput`] is returned, otherwise a `tuple` is returned where the first element is a list with the generated images. ) rKrrrr*r+r,r-rrrrMcs(g|]}t|d|dqS)rr)abs)rQrwrrr r!rTvs(z2Kandinsky5T2IPipeline.__call__..FNrr)rKrrMrLrN?z9`negative_prompt` must have same length as `prompt`. Got z vs .)rL)rrrrrNrLrr)r)rrr)totalT) rbencoder_hidden_statespooled_projectionstimestepvisual_rope_pos text_rope_pos scale_factor sparse_paramsr`)r`r)r*r+r,r-latentr)r)image)9rr r tensor_inputsrrFnpargminrrrfr.rNrrrjrgrrrr4 set_timesteps timestepsconfig in_visual_dimrrnarangerDrmaxitemorderr progress_barrhr unsqueezerrmsamplesteplocalspopupdate XLA_AVAILABLExm mark_stepr/reshapepermutescaling_factorrirE postprocessmaybe_free_model_hooksr)*rGrKrrrrrrrr)r*r+r,r-rrrr`rrrMrLrNrrrrrnegative_text_rope_posrrnum_warmup_stepsrrwtr pred_velocityuncond_pred_velocitycallback_kwargsrcallback_outputsrr rr!__call__sFN             (    6&<  zKandinsky5T2IPipeline.__call__)NrJN)NN)rrJNN)NNNNNNNN)rr8r8NNNN))__name__ __module__ __qualname____doc__model_cpu_offload_seqrrr rrrrrr@rrrnrLrrNr}rrr GeneratorTensorrpropertyrrrno_gradrEXAMPLE_DOC_STRINGfloatboolrr r r __classcell__r r rHr!r(xs>  G  '  a \  2         r()6rtypingrnumpyrregexr$rntorch.nnrrq transformersrrrr callbacksr r rEr loadersr modelsr models.transformersr schedulersrutilsrrrrutils.torch_utilsrpipeline_utilsrpipeline_outputrtorch_xla.core.xla_modelcore xla_modelrr get_loggerrrkrrr"r&r'r(r r r r!s<