o Gif @s0ddlZddlmZmZddlZddlmZmZddlm Z ddl m Z m Z ddl mZddlmZdd lmZdd lmZdd lmZmZdd lmZd dlmZeeZdZ    d"de de de!de!fddZ"    d#de dBde#ej$BdBde%e dBde%e!dBfddZ&Gd d!d!ee e Z'dS)$N)AnyCallable) AutoTokenizerPreTrainedModel)VaeImageProcessor)FromSingleFileMixinZImageLoraLoaderMixin) AutoencoderKL)ZImageTransformer2DModel)DiffusionPipeline)FlowMatchEulerDiscreteScheduler)loggingreplace_example_docstring) randn_tensor)ZImagePipelineOutputut Examples: ```py >>> import torch >>> from diffusers import ZImagePipeline >>> pipe = ZImagePipeline.from_pretrained("Z-a-o/Z-Image-Turbo", torch_dtype=torch.bfloat16) >>> pipe.to("cuda") >>> # Optionally, set the attention backend to flash-attn 2 or 3, default is SDPA in PyTorch. >>> # (1) Use flash attention 2 >>> # pipe.transformer.set_attention_backend("flash") >>> # (2) Use flash attention 3 >>> # pipe.transformer.set_attention_backend("_flash_3") >>> prompt = "一幅为名为“造相「Z-IMAGE-TURBO」”的项目设计的创意海报。画面巧妙地将文字概念视觉化：一辆复古蒸汽小火车化身为巨大的拉链头，正拉开厚厚的冬日积雪，展露出一个生机盎然的春天。" >>> image = pipe( ... prompt, ... height=1024, ... width=1024, ... num_inference_steps=9, ... guidance_scale=0.0, ... generator=torch.Generator("cuda").manual_seed(42), ... ).images[0] >>> image.save("zimage.png") ``` ?ffffff? base_seq_len max_seq_len base_shift max_shiftcCs,||||}|||}|||}|SN) image_seq_lenrrrrmbmurr`/home/ubuntu/.local/lib/python3.10/site-packages/diffusers/pipelines/z_image/pipeline_z_image.pycalculate_shift@s  r"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 valuesr%zThe current scheduler class zx's `set_timesteps` does not support custom timestep schedules. Please check whether you are using the correct scheduler.)r%r$r&zv's `set_timesteps` does not support custom sigmas schedules. Please check whether you are using the correct scheduler.)r&r$r$r) ValueErrorsetinspect signature set_timesteps parameterskeys __class__r%len) schedulerr#r$r%r&kwargsaccepts_timesteps accept_sigmasrrr!retrieve_timestepsNs2  r4c+s8eZdZdZgZddgZdededede de f fd d Z d;de e e Bdejd Bdede e e Bd Bde ejd Bdejd BdefddZ dZImagePipelineztext_encoder->transformer->vaelatents prompt_embedsr0vae text_encoder tokenizer transformercsdt|j|||||dt|dr$|jdur$dt|jjjdnd|_t |jdd|_ dS)N)r8r9r:r0r;r8r)vae_scale_factor) super__init__register_moduleshasattrr8r/configblock_out_channelsr>rimage_processor)selfr0r8r9r:r;r.rr!r@s ,zZImagePipeline.__init__NTpromptr$do_classifier_free_guidancenegative_promptnegative_prompt_embedsmax_sequence_lengthcCst|tr|gn|}|j||||d}|rB|dur!dd|D}n t|tr)|gn|}t|t|ks5J|j||||d}||fSg}||fS)N)rIr$r7rMcSsg|]}dqS)r).0_rrr! sz0ZImagePipeline.encode_prompt..) isinstancestr_encode_promptr/)rFrIr$rJrKr7rLrMrrr! encode_prompts* zZImagePipeline.encode_promptreturnc Cs|p|j}|dur |St|tr|g}t|D]\}}d|dg}|jj|dddd}|||<q|j|d|ddd}|j|} |j| } |j | | dd j d }g} t t |D] }| ||| |q[| S) Nuser)rolecontentFT)tokenizeadd_generation_promptenable_thinking max_lengthpt)paddingr] truncationreturn_tensors) input_idsattention_maskoutput_hidden_states)_execution_devicerRrS enumerater:apply_chat_templaterbtorcboolr9 hidden_statesranger/append) rFrIr$r7rMi prompt_itemmessages text_inputstext_input_ids prompt_masksembeddings_listrrr!rTsF    zZImagePipeline._encode_promptc Csdt||jd}dt||jd}||||f} |dur*t| |||d}|S|j| kr:td|jd| ||}|S)Nr<) generatorr$dtypezUnexpected latents shape, got z , expected )intr>rshaper'ri) rF batch_sizenum_channels_latentsheightwidthrvr$rur6rxrrr!prepare_latentss    zZImagePipeline.prepare_latentscC|jSr_guidance_scalerFrrr!guidance_scalezZImagePipeline.guidance_scalecCs |jdkS)Nrrrrrr!rJs z*ZImagePipeline.do_classifier_free_guidancecCr~r)_joint_attention_kwargsrrrr!joint_attention_kwargsrz%ZImagePipeline.joint_attention_kwargscCr~r)_num_timestepsrrrr! num_timestepsrzZImagePipeline.num_timestepscCr~r) _interruptrrrr! interrupt!rzZImagePipeline.interrupt2g@Fg?rpilr{r|r#r&rcfg_normalizationcfg_truncationnum_images_per_promptru output_type return_dictrcallback_on_step_end"callback_on_step_end_tensor_inputsc: s|pd}|pd}|jd}||dkr!td|d|d|d||dkr5td|d|d |d|j}||_||_d |_||_||_|d urSt|t rSd }n|d urat|t rat |}nt | }| d ury|d ury|j rx|d urxtd n|j || |j | |||d\} }|jj}|||||tj|| | } d krfdd| D} |j r|rfdd|D}|}| jdd| jdd}t||jjdd|jjdd|jjdd|jjdd}d|j_d|i}t|j||fd|i|\}}tt |||jjd}t ||_|j|dn}t|D]`\} }!|jr(q|! | jd}"d|"d}"|"d!}#|j"}$|j rY|jd urYt#|jd krY|#|jkrYd}$|j o`|$dk}%|%r}| $|jj%}&|&&dd d d }'| |}(|"&d})n | $|jj%}'| }(|"})|''d}'t |'j(dd }*|j|*|)|(d d!d}+|%r|+d |},|+|d }-g}.t)|D]F}/|,|/#}0|-|/#}1|0|$|0|1}2|jrt#|jdkrtj*+|0}3tj*+|2}4|3t#|j}5|4|5kr|2|5|4}2|.,|2qtj-|.dd }.n tj-d"d|+Ddd }.|..d}.|. }.|jj/|.$tj|!| d d!d} | j%tjks1J|d ur^i}6|D] }7t0|7|6|7<q:||| |!|6}8|81d#| } |81d$| } |81d%|}| t |d ksy| d |kr}| d |jjdkr}|2qWd n 1swY|d&kr| }9n%| $|j3j%} | |j3jj4|j3jj5} |j3j6| d d!d}9|j7j8|9|d'}9|9|s|9fSt:|9d(S))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. height (`int`, *optional*, defaults to 1024): The height in pixels of the generated image. width (`int`, *optional*, defaults to 1024): The width in pixels of the generated image. 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 to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. guidance_scale (`float`, *optional*, defaults to 5.0): Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). `guidance_scale` is defined as `w` of equation 2. of [Imagen Paper](https://arxiv.org/pdf/2205.11487.pdf). 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. cfg_normalization (`bool`, *optional*, defaults to False): Whether to apply configuration normalization. cfg_truncation (`float`, *optional*, defaults to 1.0): The truncation value for configuration. 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. 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 (`list[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. negative_prompt_embeds (`list[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. 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.ZImagePipelineOutput`] instead of a plain tuple. joint_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`, *optional*, defaults to 512): Maximum sequence length to use with the `prompt`. Examples: Returns: [`~pipelines.z_image.ZImagePipelineOutput`] or `tuple`: [`~pipelines.z_image.ZImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated images. ir<rzHeight must be divisible by z (got z-). Please adjust the height to a multiple of .zWidth must be divisible by z,). Please adjust the width to a multiple of FNrzWhen `prompt_embeds` is provided without `prompt`, `negative_prompt_embeds` must also be provided for classifier-free guidance.)rIrKrJr7rLr$rMcg|] }tD]}|qqSrrl)rOperPrrr!rQz+ZImagePipeline.__call__..crrr)rOnperPrrr!rQrrbase_image_seq_lenrmax_image_seq_lenrrrrrgr r&)totali)dim)rcSsg|]}|qSr)float)rOtrrr!rQ+sr6r7rLlatent)r)images);r>r'rfrrr_cfg_normalization_cfg_truncationrRrSlistr/rJrUr; in_channelsr}torchfloat32rxr"r0rCget sigma_minr4maxorderr progress_barrgrexpanditemrrrirvrepeat unsqueezeunbindrllinalg vector_normrmstacksqueezesteplocalspopupdater8scaling_factor shift_factordecoderE postprocessmaybe_free_model_hooksr):rFrIr{r|r#r&rrrrKrrur6r7rLrrrrrrM vae_scaler$ryrzactual_batch_sizerr scheduler_kwargsr%num_warmup_stepsrrnrtimestept_normcurrent_guidance_scale apply_cfg latents_typedlatent_model_inputprompt_embeds_model_inputtimestep_model_inputlatent_model_input_listmodel_out_listpos_outneg_out noise_predjposnegpred ori_pos_norm new_pos_norm max_new_normcallback_kwargskcallback_outputsimagerrr!__call__%s6f                              6 X zZImagePipeline.__call__)NTNNNrH)NNrHr)&__name__ __module__ __qualname__model_cpu_offload_seq_optional_components_callback_tensor_inputsr r rrr r@rSrrr$rj FloatTensorrwrUrTr}propertyrrJrrrno_gradrEXAMPLE_DOC_STRINGr Generatordictrrr __classcell__rrrGr!r5s   %   <               r5)rrrr)NNNN)(r)typingrrr transformersrrrErloadersrr models.autoencodersr models.transformersr pipelines.pipeline_utilsr schedulersr utilsrrutils.torch_utilsrpipeline_outputr get_loggerrloggerrrwrr"rSr$rr4r5rrrr!sT              ;