o
    *j	                     @   s   d dl mZmZ d dlZd dlZd dlmZ d dlm	Z	 d dl
mZ d dlmZmZ d dlmZ d dlmZ d d	lmZ d d
lmZ e ZejejejdG dd deZdS )    )AnyDictN)	Pipelines)try_on_infer)
OutputKeys)InputPipeline)	PIPELINES)	LoadImage)Tasks)
get_logger)module_namec                       s   e Zd ZdZdef fddZdedeeef fddZ	deeef deeef fd	d
Z
deeef deeef fddZ  ZS )SALForImageTryOnPipelinea   Image Try On Pipeline.
    Examples:
    >>> image_try_on = pipeline(Tasks.image_try_on, model='damo/cv_SAL-VTON_virtual-try-on', revision='v1.0.1')
    >>> input_images = {'person_input_path': '/your_path/image_try_on_person.jpg',
    >>>                 'garment_input_path': '/your_path/image_try_on_garment.jpg',
    >>>                 'mask_input_path': '/your_path/image_try_on_mask.jpg'}
    >>> result = image_try_on(input_images)
    >>> result[OutputKeys.OUTPUT_IMG]
    modelc                    sX   t  jdd|i| || _td tj r"d| _td dS d| _td dS )	z
        use `model` to create image try on pipeline for prediction
        Args:
            model: model id on modelscope hub.
        r   zload model donecudazUse GPUcpuzUse CPUN )	super__init__
model_pathloggerinfotorchr   Zis_availabledevice)selfr   kwargs	__class__r   n/var/www/html/Deteccion_Ine/venv/lib/python3.10/site-packages/modelscope/pipelines/cv/image_try_on_pipeline.pyr   !   s   

z!SALForImageTryOnPipeline.__init__inputreturnc                 C      |S Nr   )r   r   r   r   r   
preprocess2      z#SALForImageTryOnPipeline.preprocessc                 C   s0   t | j| j|d |d |d | j}tj|iS )NZperson_input_pathZgarment_input_pathZmask_input_path)r   Zinferr   r   r   r   Z
OUTPUT_IMG)r   r   Ztry_on_imager   r   r   forward5   s   
z SALForImageTryOnPipeline.forwardinputsc                 C   r!   r"   r   )r   r&   r   r   r   postprocess=   r$   z$SALForImageTryOnPipeline.postprocess)__name__
__module____qualname____doc__strr   r   r   r   r#   r%   r'   __classcell__r   r   r   r   r      s    
"*r   )typingr   r   numpynpr   Zmodelscope.metainfor   Z!modelscope.models.cv.image_try_onr   Zmodelscope.outputsr   Zmodelscope.pipelines.baser   r   Zmodelscope.pipelines.builderr	   Zmodelscope.preprocessorsr
   Zmodelscope.utils.constantr   Zmodelscope.utils.loggerr   r   Zregister_moduleZimage_try_onr   r   r   r   r   <module>   s    