o
    i                     @   sX   d dl Z d dlmZ d dlmZmZ G dd dejZG dd deZdefd	d
Z	dS )    N)override)ComfyExtensionioc                   @   s,   e Zd Zedd ZedejfddZdS )InstructPixToPixConditioningc              	   C   s^   t jddt jdt jdt jdt jdgt jjddt jjddt jjddgd	S )
Nr   zconditioning/instructpix2pixpositivenegativevaepixels)display_namelatent)node_idcategoryinputsoutputs)r   SchemaConditioningInputVaeImageOutputLatent)cls r   5/mnt/c/Users/fbmor/ComfyUI/comfy_extras/nodes_ip2p.pydefine_schema   s   



z*InstructPixToPixConditioning.define_schemareturnc                 C   s  |j d d d }|j d d d }|j d |ks |j d |krF|j d d d }|j d d d }|d d ||| ||| d d f }||}	i }
t|	|
d< g }||fD ]#}g }|D ]}|d  }|	|d< |d |g}|| q`|| qZt|d |d |
S )N         samplesconcat_latent_imager   )shapeencodetorch
zeros_likecopyappendr   
NodeOutput)r   r   r   r	   r   xyx_offsety_offsetconcat_latent
out_latentoutconditioningctdnr   r   r   execute   s&   (
z$InstructPixToPixConditioning.executeN)__name__
__module____qualname__classmethodr   r   r'   r4   r   r   r   r   r      s
    
r   c                   @   s(   e Zd Zedeeej  fddZdS )InstructPix2PixExtensionr   c                    s   t gS N)r   )selfr   r   r   get_node_list6   s   z&InstructPix2PixExtension.get_node_listN)	r5   r6   r7   r   listtyper   	ComfyNoder<   r   r   r   r   r9   5   s    r9   r   c                      s   t  S r:   )r9   r   r   r   r   comfy_entrypoint=   s   r@   )
r#   typing_extensionsr   comfy_api.latestr   r   r?   r   r9   r@   r   r   r   r   <module>   s    .