a
    Ad                     @   s0   d dl Z G dd de jZG dd de jZdS )    Nc                   @   s    e Zd ZejejdddZdS )ModuleTypeNormnetweightsc                    s$   t  fdddD r t| S d S )Nc                 3   s   | ]}| j v V  qd S )N)w).0xr    L/var/www/html/stable-diffusion-webui/extensions-builtin/Lora/network_norm.py	<genexpr>       z/ModuleTypeNorm.create_module.<locals>.<genexpr>)w_normb_norm)allNetworkModuleNormselfr   r   r
   r	   r   create_module   s    
zModuleTypeNorm.create_moduleN)__name__
__module____qualname__networkNetworkNetworkWeightsr   r
   r
   r
   r   r      s   r   c                       s0   e Zd Zejejd fddZdd Z  ZS )r   r   c                    s.   t  || |jd| _|jd| _d S )Nr   r   )super__init__r   getr   r   r   	__class__r
   r   r      s    zNetworkModuleNorm.__init__c                 C   sP   | j j}| j j|j|jd}| jd ur<| jj|j|jd}nd }| ||||S )N)dtype)r   shapetodevicer    r   finalize_updown)r   orig_weightoutput_shapeupdownex_biasr
   r
   r   calc_updown   s    
zNetworkModuleNorm.calc_updown)	r   r   r   r   r   r   r   r)   __classcell__r
   r
   r   r   r      s   r   )r   
ModuleTyper   NetworkModuler   r
   r
   r
   r   <module>   s   