o Gi@sddlmZmZddlZddlZddlmZddlm Z ddl m Z ddl m Z ddlmZdd lmZdd lmZdd lmZmZdd lmZdd lmZmZddlmZmZmZm Z m!Z!m"Z"ddl#m$Z$eruddl%m&m'Z(dZ)ndZ)e*e+Z,dZ-Gdddee Z.dS))AnyCallableN) AutoTokenizer)SmolLM3ForCausalLM)VaeImageProcessor)FluxLoraLoaderMixin)AutoencoderKLWan)BriaFiboTransformer2DModel)BriaFiboPipelineOutput)calculate_shiftretrieve_timesteps)DiffusionPipeline)FlowMatchEulerDiscreteSchedulerKarrasDiffusionSchedulers)USE_PEFT_BACKENDis_torch_xla_availableloggingreplace_example_docstringscale_lora_layersunscale_lora_layers) randn_tensorTFa Example: ```python import torch from diffusers import BriaFiboPipeline from diffusers.modular_pipelines import ModularPipeline torch.set_grad_enabled(False) vlm_pipe = ModularPipeline.from_pretrained("briaai/FIBO-VLM-prompt-to-JSON", trust_remote_code=True) pipe = BriaFiboPipeline.from_pretrained( "briaai/FIBO", trust_remote_code=True, torch_dtype=torch.bfloat16, ) pipe.enable_model_cpu_offload() with torch.inference_mode(): # 1. Create a prompt to generate an initial image output = vlm_pipe(prompt="a beautiful dog") json_prompt_generate = output.values["json_prompt"] # Generate the image from the structured json prompt results_generate = pipe(prompt=json_prompt_generate, num_inference_steps=50, guidance_scale=5) results_generate.images[0].save("image_generate.png") ``` c'@seZdZdZdZddgZdedeeBde de d e f d d Z  dHde ee BdededejdBdejdBf ddZedIddZ       dJde ee BdejdBdedede ee BdBdejdBdejdBdededBfddZed d!Zed"d#Zed$d%Zed&d'Zed(d)Zed*d+Zed,d-Z ed.d/Z!ed0d1Z"  2dKd3d4Z#ed5d6Z$e%e&e'dddd7dddd ddddd8d9dddgdd2fde ee Bd:edBd;edBdej(eej(BdBdejdBdejdBdejdBd?e dBd@e)dAe*e e+fdBdBe,eegdfdBdCee def$dDdEZ-     dLdFdGZ.dS)MBriaFiboPipelinea Args: transformer (`BriaFiboTransformer2DModel`): The transformer model for 2D diffusion modeling. scheduler (`FlowMatchEulerDiscreteScheduler` or `KarrasDiffusionSchedulers`): Scheduler to be used with `transformer` to denoise the encoded latents. vae (`AutoencoderKLWan`): Variational Auto-Encoder for encoding and decoding images to and from latent representations. text_encoder (`SmolLM3ForCausalLM`): Text encoder for processing input prompts. tokenizer (`AutoTokenizer`): Tokenizer used for processing the input text prompts for the text_encoder. z=text_encoder->text_encoder_2->image_encoder->transformer->vaelatents prompt_embeds transformer schedulervae text_encoder tokenizercCs6|j|||||dd|_t|jdd|_d|_dS)N)rrrrr)vae_scale_factor@)register_modulesr"rimage_processordefault_sample_size)selfrrrrrr(d/home/ubuntu/.local/lib/python3.10/site-packages/diffusers/pipelines/bria_fibo/pipeline_bria_fibo.py__init__^s zBriaFiboPipeline.__init__Npromptnum_images_per_promptmax_sequence_lengthdevicedtypecsp|j|p |jj}t|tr|gn|}|stdt|}d}dur'ntd}t|tjs7t|}t dd|DrRtj |df|tj |d} t | } n6|j |d|d d d d } | j|} | j|} td d|Drtjd d|Dtj|d} || | <d| | <|j| | d d} | j}tj|d|dgdd}|j|d}|jdd}tfdd|D}| jddjd} ||| fS)Nz7`prompt` must be a non-empty string or list of strings.icpucs|]}|dkVqdSNr(.0pr(r(r) z5BriaFiboPipeline.get_prompt_embeds..r+r1r0longestTpt)padding max_length truncationadd_special_tokensreturn_tensorscsr3r4r(r6r(r(r)r9r:cSsg|]}|dkqS)r5r(r6r(r(r) sz6BriaFiboPipeline.get_prompt_embeds..)attention_maskoutput_hidden_statesdimr0r1rc3s&|]}|jddjdVqdS)rrHr0N)repeat_interleavetor7layerr0r.r(r)r9s rK)_execution_devicerr1 isinstancestr ValueErrorlentorchr0allfulllong ones_liker input_idsrMrDanytensorbool hidden_statescatrLtuple)r'r-r.r/r0r1 batch_size bot_token_idtext_encoder_devicer[rD tokenized empty_rowsencoder_outputsr_rr(rPr)get_prompt_embedsrsR        z"BriaFiboPipeline.get_prompt_embedsc Cs|j\}}}|durtj||f|j|jd}n |j|j|jd}||kr(td||kr[||}tj|||f|j|jd}tj||gdd}tj||f|j|jd}tj||gdd}||fS)Nr;rJzE`max_tokens` must be greater or equal to the current sequence length.r+rH) shaperVonesr1r0rMrTzerosr`) r max_tokensrDrbseq_lenrI pad_lengthr> mask_paddingr(r(r) pad_embeddings" zBriaFiboPipeline.pad_embedding guidance_scalenegative_promptnegative_prompt_embeds lora_scalec  s|pj}| durttr| _jdurtrtj| t|tr&|gn|}|dur1t|} n|j d} d} d} |dur[j ||||d\}} } |j j j d}fdd| D} |dkrt|trl|ddurld}|pod}t|trz| |gn|}|durt|t|urtd t|d t|d | t|krtd |d t|d|d | d j ||||d\}}} |j j j d}fdd|D}jdurttrtrtj| | dur| j |j|j d} |dur2| dur| j |j|j d} t|j d|j dj|| d\}} fdd| D} j|| d\}} fdd|D}n|j dj|| d\}} d}jj }t|j ddj ||d}|||| | | |fS)a 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 guidance_scale (`float`): Guidance scale for classifier free guidance. 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.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 (`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. Nr)r-r.r/r0r1cg|] }|jjjdqSrwrMrr1r7r]r'r(r)rCz2BriaFiboPipeline.encode_prompt..r+r5z?`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`.crxryrzr{r|r(r)rC r}rJ)rDcg|] }|dqSrrprNrlr'r(r)rC5r}crrrrNrr(r)rC:r}r)rQrRr _lora_scalerrrrSrUrirhrMrr1listtype TypeErrorrTrr0maxrprVrk)r'r-r0r.rsrtrrur/rvrbprompt_attention_masknegative_prompt_attention_mask prompt_layersnegative_prompt_layersr1text_idsr(rr) encode_prompts "              zBriaFiboPipeline.encode_promptcC|jSN)_guidance_scaler|r(r(r)rsOzBriaFiboPipeline.guidance_scalecCrr)_joint_attention_kwargsr|r(r(r)joint_attention_kwargsWrz'BriaFiboPipeline.joint_attention_kwargscCrr)_num_timestepsr|r(r(r) num_timesteps[rzBriaFiboPipeline.num_timestepscCrr) _interruptr|r(r(r) interrupt_rzBriaFiboPipeline.interruptcCsh|j\}}}||}||}|||d|d|ddd}|dddddd}|||d||}|S)Nr!rrr+rq)riviewpermutereshaperheightwidthr"rb num_patcheschannelsr(r(r)_unpack_latentscs  z BriaFiboPipeline._unpack_latentsc Cs|t||d}|dt|dddf|d<|dt|dddf|d<|j\}}}||||}|j||dS)Nr).r+).r!rJ)rVrkarangerirrM) rbrrr0r1latent_image_idslatent_image_id_heightlatent_image_id_widthlatent_image_id_channelsr(r(r)_prepare_latent_image_idsqs"" z*BriaFiboPipeline._prepare_latent_image_idscCs@|j\}}}||}||}|||||}|dddd}|S)Nrrr+r!)rirrrr(r(r)_unpack_latents_no_patchs z)BriaFiboPipeline._unpack_latents_no_patchcCs&|dddd}|||||}|S)Nrr!rr+)rrrrbnum_channels_latentsrrr(r(r)_pack_latents_no_patchsz'BriaFiboPipeline._pack_latents_no_patchcCsR||||dd|dd}|dddddd}|||d|d|d}|S)Nr!rrr+rrq)rrrrr(r(r) _pack_latentsszBriaFiboPipeline._pack_latentsFc Cst||j}t||j}||||f} |dur*||||||} |j||d| fSt|trBt||krBtdt|d|dt| |||d}| rf| |||||}|||d|d||} || fS| |||||}||||||} || fS)NrJz/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.) generatorr0r1r!) intr"rrMrRrrUrTrrr) r'rbrrrr1r0rr do_patchingrirr(r(r)prepare_latentss(   z BriaFiboPipeline.prepare_latentscCs(td||}t|dkdtj }|S)Nz bi,bj->bijr+g)rVeinsumwhereinf)rDattention_matrixr(r(r)_prepare_attention_masks z(BriaFiboPipeline._prepare_attention_maskpilTrrnum_inference_steps timestepsr output_type return_dictrcallback_on_step_end"callback_on_step_end_tensor_inputsc5 s@|p|j|j}|p|j|j}|j|||| ||d||_||_d|_|dur0t|tr0d}n|dur>t|tr>t |}n| j d}|j }|j durR|j ddnd}|j|||| | ||||d \} } }}}| j d}|dkrtj| | gdd} fd d tt Dtj||gdd}t |jjt |jj}t |krt |dn d g|t |jjj}|rt|d }|||||| j|| | | \} }tj| j d| j dg| j| jd }|dkr|dd}tj||gdd}||}|jddj|jjd}|jduri|_||jd<|r/||jd||jd} n ||j||j} t !dd||}!t"| |j#jj$|j#jj%|j#jj&|j#jj'}"t(|j#||d|!|"d\}}t)t |||j#j*d}#t ||_+t |j dkr|d}t |j dkr|d}|j,|d}$t-|D]\}%}&|j.rq|dkrt| gdn| }'|&/|'j dj|'j|'jd}(|j|'|(| |j d||dd})|dkr|)0d\}*}+|*|j1|+|*})| j},|j#j2|)|&| ddd} | j|,krtj3j45r| |,} |dur3i}-|D] }.t6|.|-|.<q|||%|&|-}/|/7d| } |/7d| } |/7d| } |%t |dksN|%d|#krR|%d|j#j*dkrR|$8t9rYt:;qWdn 1sfwY| dkrs| }0n|r|<| |||j} n |=| |||j} | jdd} | dj}1| dj},t>|j?jj@Ad|j?jjBddd|1|,dt>|j?jjCAd|j?jjBddd|1|,fdd | D}2tj|2dd}2g}0|2D]!}3|j?jD|3dddd}4|jEjF|4jGdd| d}4|0H|4qt |0dkr |0d}0nt jI|0dd}0|J|s|0fStK|0dS) aH 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 self.unet.config.sample_size * self.vae_scale_factor): The height in pixels of the generated image. This is set to 1024 by default for the best results. width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): The width in pixels of the generated image. This is set to 1024 by default for the best results. 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. timesteps (`list[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. guidance_scale (`float`, *optional*, defaults to 5.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. 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 ge generated by sampling using the supplied random `generator`. prompt_embeds (`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 (`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_xl.StableDiffusionXLPipelineOutput`] 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` defaults to 3000): Maximum sequence length to use with the `prompt`. do_patching (`bool`, *optional*, defaults to `False`): Whether to use patching. Examples: Returns: [`~pipelines.flux.BriaFiboPipelineOutput`] or `tuple`: [`~pipelines.flux.BriaFiboPipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated images. )r-rrrrr/FNr+rscale) r-rtrsrrur0r/r.rvrHcs&g|]}tj||gddqS)rrH)rVr`)r7i)rrr(r)rC`sz-BriaFiboPipeline.__call__..rFrr;r!rwrDg?)rr0rsigmasmur)totalrJ)r_timestepencoder_hidden_statestext_encoder_layersrrtxt_idsimg_ids)rrrrulatentcsg|]}|qSr(r()r7r) latents_mean latents_stdr(r)rCs)r)axis)images)Lr&r" check_inputsrrrrRrSrrUrirQrgetrrVr`rangertransformer_blockssingle_transformer_blocksconfig in_channelsrrr1rjr0repeatr unsqueezerMnplinspacer rbase_image_seq_lenmax_image_seq_len base_shift max_shiftr rorderr progress_bar enumeraterexpandchunkrsstepbackendsmps is_availablelocalspopupdate XLA_AVAILABLExm mark_steprrr]rrrz_dimrdecoder% postprocesssqueezeappendstackmaybe_free_model_hooksr )5r'r-rrrrrsrtr.rrrrurrrrrr/rrbr0rvrrrprompt_batch_sizetotal_num_layers_transformerrrlatent_attention_maskrDrmrrnum_warmup_stepsrrtlatent_model_inputr noise_prednoise_pred_uncondnoise_pred_text latents_dtypecallback_kwargskcallback_outputsimagelatents_devicelatents_scaled scaled_latent curr_imager()rrrrr)__call__s^`                           6 8   &  zBriaFiboPipeline.__call__c s^|ddks |ddkrtd|d|d|dur8tfdd|Ds8tdjd fd d |D|durK|durKtd |d |d|durW|durWtd|durnt|tsnt|tsntdt||dur|durtd|d|d|dur|dur|j|jkrtd|jd|jd|dur|dkrtd|dSdS)Nr rz8`height` and `width` have to be divisible by 16 but are z and r~c3s|]}|jvVqdSr_callback_tensor_inputsr7rr|r(r)r9$s  z0BriaFiboPipeline.check_inputs..z2`callback_on_step_end_tensor_inputs` has to be in z , but found csg|] }|jvr|qSr(rrr|r(r)rC(sz1BriaFiboPipeline.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` rrz9`max_sequence_length` cannot be greater than 3000 but is )rTrWrrRrSrrri) r'r-rrrtrrurr/r(r|r)rsH  zBriaFiboPipeline.check_inputs)r+r,NNr)Nr+rqNNNrrN)NF)NNNNN)/__name__ __module__ __qualname____doc__model_cpu_offload_seqrr rrr rrr*rSrrrVr0r1rh staticmethodrpfloat FloatTensorrpropertyrsrrrrrrrrrrno_gradrEXAMPLE_DOC_STRING Generatorr^dictrrrrr(r(r(r)rLs*   <              %        Qr)/typingrrnumpyrrV transformersr,transformers.models.smollm3.modeling_smollm3rr%rloadersr&models.autoencoders.autoencoder_kl_wanr )models.transformers.transformer_bria_fibor #pipelines.bria_fibo.pipeline_outputr pipelines.flux.pipeline_fluxr r pipelines.pipeline_utilsr schedulersrrutilsrrrrrrutils.torch_utilsrtorch_xla.core.xla_modelcore xla_modelrr get_loggerrloggerrrr(r(r(r)s,