o Gi @sPddlZddlZddlZddlmZddlmZddl m Z GdddeZ dS)N)register_to_config)VaeImageProcessor)PIL_INTERPOLATIONcseZdZdZe           d d ed ed ed eeefdededededededee Bee dfBdBffdd Z de j j dedede j j fddZ  d!de j j ejBejBdedBdedBdeeeffddZZS)"WanAnimateImageProcessorat Image processor to preprocess the reference (character) image for the Wan Animate 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 `8`): 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 `16`): VAE latent channels. spatial_patch_size (`tuple[int, int]`, *optional*, defaults to `(2, 2)`): The spatial patch size used by the diffusion transformer. For Wan models, this is typically (2, 2). resample (`str`, *optional*, defaults to `lanczos`): Resampling filter to use when resizing the image. do_normalize (`bool`, *optional*, defaults to `True`): Whether to normalize the image to [-1,1]. do_binarize (`bool`, *optional*, defaults to `False`): Whether to binarize the image to 0/1. do_convert_rgb (`bool`, *optional*, defaults to be `False`): Whether to convert the images to RGB format. do_convert_grayscale (`bool`, *optional*, defaults to be `False`): Whether to convert the images to grayscale format. fill_color (`str` or `float` or `tuple[float, ...]`, *optional*, defaults to `None`): An optional fill color when `resize_mode` is set to `"fill"`. This will fill the empty space with that color instead of filling with data from the image. Any valid `color` argument to `PIL.Image.new` is valid; if `None`, will default to filling with data from `image`. Tr lanczosNFr do_resizevae_scale_factorvae_latent_channelsspatial_patch_sizeresample reducing_gap do_normalize do_binarizedo_convert_rgbdo_convert_grayscale fill_color.c s$t| r| rtdddSdS)Nz`do_convert_rgb` and `do_convert_grayscale` can not both be set to `True`, if you intended to convert the image into RGB format, please set `do_convert_grayscale = False`.z` if you intended to convert the image into grayscale format, please set `do_convert_rgb = False`)super__init__ ValueError) selfr r rrrrrrrrr __class__[/home/ubuntu/.local/lib/python3.10/site-packages/diffusers/pipelines/wan/image_processor.pyr8s z!WanAnimateImageProcessor.__init__imagewidthheightreturncCs||}|j|j}|jjdu}|jjpd}||kr|n|j||j}||kr*|n|j||j} |j|| ft|jjd} tjj d||f|d} | j | |d|d|d| dfd|r||kr|d| d} | dkr| j | j|| fdd|dfddd| j | j|| fd| j|| jfdd| | fd| S||kr|d|d} | dkr| j | j| |fddd|fddd| j | j| |f| jd| j|fd| |dfd| S) a* Resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center the image within the dimensions, filling empty with data from image. Args: image (`PIL.Image.Image`): The image to resize and fill. 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 filled image. Nr)rRGB)colorr )box)rr) r r!configrresizerrPILImagenewpaste)rrr r!ratio src_ratiofill_with_image_datarsrc_wsrc_hresizedres fill_height fill_widthrrr_resize_and_fillOs8   *$  $ z)WanAnimateImageProcessor._resize_and_fillcCs|dur t|tjjr|j}nt|tjr|jd}n|jd}|dur@t|tjjr/|j}nt|tjr;|jd}n|jd}||}||}|jj |jj d}|jj |jj d}t t ||||}t t ||||}||fS)a= Returns the height and width of the image, downscaled to the next integer multiple of `vae_scale_factor`. Args: image (`PIL.Image.Image | np.ndarray | torch.Tensor`): The image input, which can be a PIL image, NumPy array, or PyTorch tensor. If it is a NumPy array, it should have shape `[batch, height, width]` or `[batch, height, width, channels]`. If it is a PyTorch tensor, it should have shape `[batch, channels, height, width]`. height (`int | None`, *optional*, defaults to `None`): The height of the preprocessed image. If `None`, the height of the `image` input will be used. width (`int | None`, *optional*, defaults to `None`): The width of the preprocessed image. If `None`, the width of the `image` input will be used. Returns: `tuple[int, int]`: A tuple containing the height and width, both resized to the nearest integer multiple of `vae_scale_factor * spatial_patch_size`. Nr rr) isinstancer(r)r!torchTensorshaper r&r rroundnpsqrt)rrr!r max_area aspect_ratio mod_value_h mod_value_wrrrget_default_height_widths&      z1WanAnimateImageProcessor.get_default_height_width) Trrr r NTFFFr)NN)__name__ __module__ __qualname____doc__rboolinttuplestrfloatrr(r)r5r<ndarrayr8r9rB __classcell__rrrrrsr      : r) numpyr< PIL.Imager(r8configuration_utilsrimage_processorrutilsrrrrrrs