o
    ۷i                     @   s<   d dl Z d dlZddlmZ ddlmZ G dd deZdS )    N   )register_to_config)VaeImageProcessorc                       s   e Zd ZdZe					d"dedededed	ef
 fd
dZe	d#de	j
j
dededede	j
j
f
ddZed$de	j
j
dede	j
j
fddZed$de	j
j
fddZde	j
j
dedede	j
j
fddZedee	j
j
 de	j
j
fd d!Z  ZS )%Flux2ImageProcessora  
    Image processor to preprocess the reference (character) image for the Flux2 model.

    Args:
        do_resize (`bool`, *optional*, defaults to `True`):
            Whether to downscale the image's (height, width) dimensions to multiples of `vae_scale_factor`. Can accept
            `height` and `width` arguments from [`image_processor.VaeImageProcessor.preprocess`] method.
        vae_scale_factor (`int`, *optional*, defaults to `16`):
            VAE (spatial) scale factor. If `do_resize` is `True`, the image is automatically resized to multiples of
            this factor.
        vae_latent_channels (`int`, *optional*, defaults to `32`):
            VAE latent channels.
        do_normalize (`bool`, *optional*, defaults to `True`):
            Whether to normalize the image to [-1,1].
        do_convert_rgb (`bool`, *optional*, defaults to be `True`):
            Whether to convert the images to RGB format.
    T       	do_resizevae_scale_factorvae_latent_channelsdo_normalizedo_convert_rgbc                    s   t  j|||||d d S )N)r   r	   r
   r   r   )super__init__)selfr   r	   r
   r   r   	__class__ _/home/ubuntu/vllm_env/lib/python3.10/site-packages/diffusers/pipelines/flux2/image_processor.pyr   *   s   	
zFlux2ImageProcessor.__init__   @      imagemax_aspect_ratiomin_side_lengthmax_areareturnc              
   C   s   t | tjjstdt|  | j\}}||k s||k r+td| d| d| dt|| || }||krJtd| d| d|dd	| d
	| S )u  
        Check if image meets minimum size and aspect ratio requirements.

        Args:
            image: PIL Image to validate
            max_aspect_ratio: Maximum allowed aspect ratio (width/height or height/width)
            min_side_length: Minimum pixels required for width and height
            max_area: Maximum allowed area in pixels²

        Returns:
            The input image if valid

        Raises:
            ValueError: If image is too small or aspect ratio is too extreme
        z%Image must be a PIL.Image.Image, got zImage too small:    ×z#. Both dimensions must be at least pxzAspect ratio too extreme: z	 (ratio: z.1fz:1). Maximum allowed ratio is z:1)
isinstancePILImage
ValueErrortypesizemax)r   r   r   r   widthheightaspect_ratior   r   r   check_image_input;   s   
z%Flux2ImageProcessor.check_image_inputtarget_areac                 C   sJ   | j \}}t|||  }t|| }t|| }| ||ftjjjS N)	r#   mathsqrtintresizer   r    
ResamplingLANCZOS)r   r)   image_widthimage_heightscaler%   r&   r   r   r   _resize_to_target_areac   s
   
z*Flux2ImageProcessor._resize_to_target_areac                 C   s*   | j \}}|| }||kr| S t| |S r*   )r#   r   r4   )r   r)   r1   r2   pixel_countr   r   r   _resize_if_exceeds_aream   s
   
z+Flux2ImageProcessor._resize_if_exceeds_arear%   r&   c           
      C   sD   |j \}}|| d }|| d }|| }|| }	|||||	fS )a  
        center crop the image to the specified width and height.

        Args:
            image (`PIL.Image.Image`):
                The image to resize and crop.
            width (`int`):
                The width to resize the image to.
            height (`int`):
                The height to resize the image to.

        Returns:
            `PIL.Image.Image`:
                The resized and cropped image.
           )r#   crop)
r   r   r%   r&   r1   r2   lefttoprightbottomr   r   r   _resize_and_cropu   s   
z$Flux2ImageProcessor._resize_and_cropimagesc                 C   s   t | dkr| d  S dd | D } tdd | D }tdd | D }d}tjd	||f|}d}| D ]}||j d
 }||||f ||j	7 }q5|S )zk
        Concatenate a list of PIL images horizontally with center alignment and white background.
           r   c                 S   s$   g | ]}|j d kr|d n|qS )RGB)modeconvert.0imgr   r   r   
<listcomp>   s   $ z:Flux2ImageProcessor.concatenate_images.<locals>.<listcomp>c                 s       | ]}|j V  qd S r*   )r%   rC   r   r   r   	<genexpr>       z9Flux2ImageProcessor.concatenate_images.<locals>.<genexpr>c                 s   rG   r*   )r&   rC   r   r   r   rH      rI   )   rJ   rJ   r@   r7   )
lencopysumr$   r   r    newr&   paster%   )r>   total_width
max_heightbackground_colornew_imgx_offsetrE   y_offsetr   r   r   concatenate_images   s   z&Flux2ImageProcessor.concatenate_images)Tr   r   TT)r   r   r   )r   )__name__
__module____qualname____doc__r   boolr-   r   staticmethodr   r    r(   r4   r6   r=   listrV   __classcell__r   r   r   r   r      s^    '"	
 (r   )r+   	PIL.Imager   configuration_utilsr   image_processorr   r   r   r   r   r   <module>   s
   