
    +jE                     |   d dl Z ddlmZ ej        Zd Zee j        e j        gej        _        e j	        ej        _
        e j	        e j        gej        _        e j        ej        _
        e j	        gej        _        e j        ej        _
        e j	        gej        _        e j        ej        _
         G d d          ZdS )    N   )controlc                  |    t          t          t          dd                                                    } d| v rdndS )N_name rocmi    i   )strgetattrliblower)lib_names    Z/home/wildlama/comfy/ComfyUI/.venv/lib/python3.11/site-packages/comfy_aimdo/vram_buffer.py_growth_chunk_sizer      s;    73,,--3355H H,,;;,>    c                   (    e Zd Zd Zd ZddZd ZdS )
VRAMBufferc                 P   t          j        |          | _        || _        || _        t
                              | j        | j                  | _        | j        st          d          t
          	                    | j                  | _
        d| _        t                      | _        d S )NzVRAM reservation failedr   )r   
get_devctx_devctxdevicemax_sizer   vrambuf_create_ptrRuntimeErrorvrambuf_get	base_addr
_allocatedr   _chunk_size)selfr   r   s      r   __init__zVRAMBuffer.__init__   s    )&11 &&t{DMBB	y 	:899933-//r   c                     | j         S )N)r   )r   s    r   sizezVRAMBuffer.size(   s
    r   r   c                    t          |          }t          |          }||z   }|| j        k    rQt                              | j        |          st          d| d          || j        z   dz
  | j        dz
   z  | _        | | j        |z   |fS )NzVRAM grow failed: z bytesr   )intr   r   vrambuf_growr   r   r   r   )r   r"   offsetrequired_sizes       r   getzVRAMBuffer.get+   s    V4yyv4?**##DI}== O"#M#M#M#MNNN,t/??!CIY\]I]G^^DOdnv-t44r   c                     t          | dd           }|r2t                              |          st          d          d | _        d S d S )Nr   z/VRAM destroy failed: device context unavailable)r
   r   vrambuf_destroyr   r   )r   ptrs     r   __del__zVRAMBuffer.__del__7   sV    dFD)) 	&&s++ V"#TUUUDIII	 	r   N)r   )__name__
__module____qualname__r    r"   r(   r,    r   r   r   r      sU        
0 
0 
0  
5 
5 
5 
5    r   r   )ctypesr   r   r   r   c_intc_size_tr   argtypesc_void_prestyper%   c_boolr   c_uint64r*   r   r0   r   r   <module>r9      s          k? ? ?
 ?#)<"AC!'C!'&/ BC%}C &0CO$oCO$*O#4C "(-C! ! ! ! ! ! ! ! ! !r   