a
    AdE                     @   s0  d dl Zd dlZd dlmZ d dlmZmZm	Z	m
Z
 d dlmZmZ d dlZd dlZd dlZd dlmZ d dlmZ d dlmZ g Ze Zdd	 Zd)edddZd*eedddZd+eeedddZdd Z dd Z!G dd dZ"dd Z#dd Z$G dd dZ%d d! Z&ej'd"d#d$Z(d%d& Z)d'd( Z*dS ),    N)Path)sharedui_extra_networks_user_metadataerrorsextra_networks)read_info_from_imagesave_image_with_geninfo)HTTPException)image_from_url_text)
ToolButtonc                 C   s4   t |  t  tttdd t D g  dS )zfregisters extra networks page for the UI; recommend doing it in on_before_ui() callback for extensionsc                 S   s   g | ]}|  qS  ) allowed_directories_for_previews.0xr   r   A/var/www/html/stable-diffusion-webui/modules/ui_extra_networks.py
<listcomp>       z!register_page.<locals>.<listcomp>N)extra_pagesappendallowed_dirsclearupdatesetsumpager   r   r   register_page   s    
r    filenamec                    s   ddl m} tj s$tdddt fddtD sJtd  d	tj	 d
 
 }|dvrvtd  d| ddidS )Nr   )FileResponsei  zFile not found)status_codedetailc                 3   s(   | ] }t | t   jv V  qd S N)r   absoluteparentsr   r   r   r   	<genexpr>!   r   zfetch_file.<locals>.<genexpr>zFile cannot be fetched: z:. Must be in one of directories registered by extra pages.   )z.pngz.jpgz.jpegz.webpz.gifz. Only png, jpg, webp, and gif.zAccept-Rangesbytes)headers)starlette.responsesr!   ospathisfiler	   anyr   
ValueErrorsplitextlower)r    r!   extr   r   r   
fetch_file   s    r4   )r   itemc                    sl   ddl m} tt fddtD d   d u r8|i S  j|}|d u rT|i S |dtj|dddiS )	Nr   JSONResponsec                    s   g | ]}|j  kr|qS r   namer   r   r   r   r   /   r   z get_metadata.<locals>.<listcomp>metadata   F)indentensure_ascii)	r+   r7   nextiterr   r:   getjsondumps)r   r5   r7   r:   r   r   r   get_metadata,   s    rC   )r   tabnamer9   c              
      s   ddl m} tt fddtD d  z j|dd}| j|< W n< ty } z$t	|d  j
|}W Y d }~n
d }~0 0  |  ||}|d|iS )	Nr   r6   c                    s   g | ]}|j  kr|qS r   r8   r   r   r   r   r   =   r   z#get_single_card.<locals>.<listcomp>F)Zenable_filterzcreating item for extra networkhtml)r+   r7   r>   r?   r   create_itemitems	Exceptionr   displayr@   read_user_metadatacreate_html_for_item)r   rD   r9   r7   r5   eZ	item_htmlr   r   r   get_single_card:   s    "
rM   c                 C   s:   | j dtdgd | j dtdgd | j dtdgd d S )Nz/sd_extra_networks/thumbGET)methodsz/sd_extra_networks/metadataz"/sd_extra_networks/get-single-card)add_api_router4   rC   rM   )appr   r   r   add_pages_to_demoL   s    rR   c                 C   s$   |  dd} |  dd} d|  dS )N\z\\"z\")replace)sr   r   r   quote_jsR   s    rW   c                   @   s   e Zd Zdd Zdd Zdd Zdd Zdd
dZdd ZdddZ	dd Z
dd Zdd Zdd Zdd Zdd Zdd Zd	S ) ExtraNetworksPagec                 C   sB   || _ | | _| jdd| _td| _d| _i | _	i | _
d S )N _zextra-networks-card.htmlF)titler2   r9   rU   id_pager   rE   	card_pageallow_negative_promptr:   rG   )selfr[   r   r   r   __init__Y   s    
zExtraNetworksPage.__init__c                 C   s   d S r$   r   r_   r   r   r   refreshb   s    zExtraNetworksPage.refreshc                 C   s>   | dd }t|}| dd }|d ur2||d< ||d< d S )Nr    descriptionuser_metadata)r@   r   get_user_metadata)r_   r5   r    r:   descr   r   r   rJ   e   s    
z$ExtraNetworksPage.read_user_metadatac                 C   s0   t j|dd}tj|}d| d| S )NrS   /z#./sd_extra_networks/thumb?filename=z&mtime=)urllibparsequoterU   r,   r-   getmtime)r_   r    Zquoted_filenamemtimer   r   r   link_previewo   s    zExtraNetworksPage.link_previewNc                 C   s\   t j|}|d ur|n|  D ]6}t j|}||r |t|d  dd  S q dS )NrS   rg   r   )r,   r-   abspathr   
startswithlenrU   )r_   r    Zpossible_directoriesrn   	parentdirr   r   r   search_terms_from_patht   s    
z(ExtraNetworksPage.search_terms_from_pathc                    s  d}i | _ i }dd |  D D ]}ttj|dddd dD ]\}}}t|tjdD ]}tj||}	tj	|	svqVtj
|	t|d  d	d
}
|
d
r|
dd  }
qtt|	dk}|s|
d
s|
d
 }
d|
v s|
drtjjsqVd||
< qVq>q |rddi|}d fdd|D }dd |  D | _| j D ]H}|d}|rj|| j |d < d|vr~| | || | 7 }qH|dkrddd |  D }tdj|d}| jdd}d  d| d| d  d| d| d}|S )Nr   c                 S   s   g | ]}t j|qS r   )r,   r-   rn   r   r   r   r   r      r   z1ExtraNetworksPage.create_html.<locals>.<listcomp>T)followlinksc                 S   s   t | d S )Nr   )r   natural_sort_keyr   r   r   r   <lambda>   r   z/ExtraNetworksPage.create_html.<locals>.<lambda>keyrS   rg   r(   r   /..c                    sB   g | ]:}d |dkrdnd d  dt |dkr2|nd dqS )z8
<button class='lg secondary gradio-button custom-buttonr   z search-allz&' onclick='extraNetworksSearchButton("z_extra_search", event)'>
allz
</button>
rE   escape)r   subdirrD   r   r   r      s   c                 S   s   i | ]}|d  |qS r8   r   r   r   r   r   
<dictcomp>   r   z1ExtraNetworksPage.create_html.<locals>.<dictcomp>r:   r9   rd   c                 S   s   g | ]}d | dqS )z<li>z</li>r   r   r   r   r   r      r   zextra-networks-no-cards.html)dirsrY   rZ   z

<div id='zE_subdirs' class='extra-network-subdirs extra-network-subdirs-cards'>
z
</div>
<div id='z%_cards' class='extra-network-cards'>
z
</div>
)r:   r   sortedr,   walkr   rt   r-   joinisdirrn   rp   rU   ro   listdirendswithopts&extra_networks_show_hidden_directories
list_itemsrG   valuesr@   rJ   rK   rE   formatr9   )r_   rD   Z
items_htmlsubdirsrq   rootr   rZ   dirnamer   r~   is_emptyZsubdirs_htmlr5   r:   Zself_name_idresr   r   r   create_html~   sb    $ 




	zExtraNetworksPage.create_htmlc                 C   s
   t  d S r$   NotImplementedError)r_   r9   indexr   r   r   rF      s    zExtraNetworksPage.create_itemc                 C   s
   t  d S r$   r   ra   r   r   r   r      s    zExtraNetworksPage.list_itemsc                 C   s   g S r$   r   ra   r   r   r   r      s    z2ExtraNetworksPage.allowed_directories_for_previewsc                 C   sp  | dd}| dd}|du rXdtdt| d|d  d| jrFdnd	 d
 d }tjjrpdtjj dnd}tjjrdtjj dnd}|rdt| dnd}d}| d}	|	rdt| j	 dt|d  d}dt| dt| j
 dt|d  d}
d}| dd}|  D ].}tj|}||r|t|d }qtjjdkrbd}nd|v prd|v }|rtjjdkrdS ddd | di  D  }|d | | d!tjjd"  d#| ddt|t|d$ t|d tjjr| d%pdnd|dtd&t| dt|d$  d
 d | d'd||
|rTd(nd|d)}| jjf i |S )*zy
        Create HTML for card item in tab tabname; can return empty string if the item is not meant to be shown.
        previewNonclickrT   zreturn cardClicked(z, prompttruefalse)zheight: zpx;r   zwidth: z
<img src="z!" class="preview" loading="lazy">r:   zu<div class='metadata-button card-button' title='Show internal metadata' onclick='extraNetworksRequestMetadata(event, r9   z	)'></div>zi<div class='edit-button card-button' title='Edit metadata' onclick='extraNetworksEditUserMetadata(event, r    AlwaysFry   z\.NeverrY   c                 S   s&   g | ]\}}t d | d| qS )z
data-sort-=r|   )r   kvr   r   r   r      r   z:ExtraNetworksPage.create_html_for_item.<locals>.<listcomp>	sort_keysz'display: none; z; font-size: d   z%'local_previewrc   zreturn saveCardPreview(event, search_termz search_only)background_imagestyler   rD   r   r9   rc   Zcard_clickedZsave_card_previewr   metadata_buttonedit_buttonsearch_onlyr   )r@   rE   r}   rW   r^   r   r   extra_networks_card_heightextra_networks_card_widthr9   r\   r   r,   r-   rn   ro   rp   extra_networks_hidden_modelsr   rG   stripextra_networks_card_text_scaleextra_networks_card_show_descr]   r   )r_   r5   rD   r   r   heightwidthr   r   r:   r   
local_pathr    ZreldirZabsdirr   r   argsr   r   r   rK      sP    8
 *$

*
z&ExtraNetworksPage.create_html_for_itemc                 C   s6   t |}| }t|jpdt|jp&d|j dS )zB
        List of default keys used for sorting in the UI.
        r   )Zdate_createdZdate_modifiedr9   )r   statintst_ctimest_mtimer9   r2   )r_   r-   pthr   r   r   r   get_sort_keys  s    zExtraNetworksPage.get_sort_keysc                    sb   g d}t jj|vr"|t jj t fdd|D g }|D ]}tj|r>| |  S q>dS )zf
        Find a preview PNG for a given path (without extension) and call link_preview on it.
        )pngjpgjpegwebpc                    s$   g | ]} d  |  d | gqS )rz   z	.preview.r   )r   r3   r-   r   r   r     r   z2ExtraNetworksPage.find_preview.<locals>.<listcomp>N)	r   r   samples_formatr   r   r,   r-   r.   rm   )r_   r-   Zpreview_extensionsZpotential_filesfiler   r   r   find_preview  s    zExtraNetworksPage.find_previewc              	   C   st   | d| dfD ]Z}zBt |dddd}| W  d   W   S 1 sN0    Y  W q tyl   Y q0 qdS )zX
        Find and read a description file for a given path (without extension).
        z.txtz.description.txtrzutf-8rU   )encodingr   N)openreadOSError)r_   r-   r   fr   r   r   find_description   s    0z"ExtraNetworksPage.find_descriptionc                 C   s   t ||| S r$   )r   ZUserMetadataEditor)r_   uirD   r   r   r   create_user_metadata_editor,  s    z-ExtraNetworksPage.create_user_metadata_editor)N)N)__name__
__module____qualname__r`   rb   rJ   rm   rr   r   rF   r   r   rK   r   r   r   r   r   r   r   r   rX   X   s   	


@
<rX   c                   C   s   t   d S r$   )r   r   r   r   r   r   
initialize0  s    r   c                  C   sF   ddl m}  ddlm} ddlm} t|   t|  t|  d S )Nr   )!ExtraNetworksPageTextualInversion)ExtraNetworksPageHypernetworks)ExtraNetworksPageCheckpoints)Z+modules.ui_extra_networks_textual_inversionr   Z#modules.ui_extra_networks_hypernetsr   Z%modules.ui_extra_networks_checkpointsr   r   )r   r   r   r   r   r   register_default_pages4  s    

r   c                   @   s   e Zd Zdd ZdS )ExtraNetworksUic                 C   s(   d | _ d | _d | _d | _d | _d | _d S r$   )pagespage_contentsstored_extra_pagesbutton_save_previewpreview_target_filenamerD   ra   r   r   r   r`   >  s    zExtraNetworksUi.__init__N)r   r   r   r`   r   r   r   r   r   =  s   r   c                    sP   dd t jjdD  fddfddt D t fdd	d
S )Nc                 S   s   g | ]}|   qS r   )r2   r   r   r   r   r   r   N  r   z,pages_in_preferred_order.<locals>.<listcomp>,c                    s2   |   } tD ]\}}|| v r|  S qt S r$   )r2   	enumeraterp   )r9   ipossible_match)r   	tab_orderr   r   tab_name_scoreP  s
    
z0pages_in_preferred_order.<locals>.tab_name_scorec                    s"   i | ]\}}|j  |j |fqS r   r8   )r   original_indexr   )r   r   r   r   X  r   z,pages_in_preferred_order.<locals>.<dictcomp>c                    s
    | j  S r$   r8   ru   )
tab_scoresr   r   rv   Z  r   z*pages_in_preferred_order.<locals>.<lambda>rw   )r   r   ui_extra_networks_tab_reordersplitr   r   )r   r   )r   r   r   r   r   pages_in_preferred_orderM  s    r   )	interfacec                    s@  ddl m} t g _g _g _tt _	|_
g }j	D ]}tj|j|jd}| d|j d}tjd|d}j| |jdd	 d
t| d g g d ||}	|	  j|	 || W d    qB1 s0    Y  qBtjdd|d ddddd}
tjg dd|d ddddd|d d	}t||d ddd}tjd|d dd}tjdd |d! d"dd#}tjd$|d% dd_tjd&|d' dd_|D ]&}|jd(d	 g |
||||gdd) q|D ]&}|jd*d	 g |
||||gdd) q҇ fd+d,}fd-d. | j|g g jd/ |j g jd/ S )0Nr   )switch_values_symbol)idrZ   Z_cards_htmlz
Loading...)elem_idc                   S   s   d S r$   r   r   r   r   r   rv   o  r   zcreate_ui.<locals>.<lambda>z#function(){applyExtraNetworkFilter(z); return []}fn_jsinputsoutputsr   FZ_extra_searchsearchz	Search...T)
show_labelr   elem_classesplaceholdervisibleinteractive)Default SortzDate CreatedzDate ModifiedNamer   Z_extra_sortsortZ_extra_sort_order)	choicesvaluer   r   multiselectr   r   r   labelZ_extra_sortorder	sortorder)r   r   r   RefreshZ_extra_refresh)r   r   z	Show dirsZ_extra_show_dirsz	show-dirs)r   r   r   r   zSave previewZ_save_previewzPreview save filenameZ_preview_filenamec                   S   s   dd t dD S )Nc                 S   s   g | ]}t jd dqS )Fr   grr   r   rZ   r   r   r   r     r   /create_ui.<locals>.<lambda>.<locals>.<listcomp>   ranger   r   r   r   rv     r   )r   r   r   show_progressc                   S   s   dd t dD S )Nc                 S   s   g | ]}t jd dqS )Tr   r   r   r   r   r   r     r   r   r  r  r   r   r   r   rv     r   c                      s   j s  S j S r$   )pages_contentsr   rb   r   r   r   
pages_html  s    zcreate_ui.<locals>.pages_htmlc                     s0    j D ]} |   q fdd j D  _ jS )Nc                    s   g | ]}|  jqS r   r   rD   )r   pgr   r   r   r     r   z.create_ui.<locals>.refresh.<locals>.<listcomp>)r   rb   r  )r	  r
  r   r   rb     s    

zcreate_ui.<locals>.refresh)r   r   r   )
modules.uir   r   r   r  user_metadata_editorsr   r   copyr   rD   r   Tabr[   r\   HTMLr   changerW   r   	create_uiTextboxDropdownr   ButtonCheckboxr   r   selectloadclick)r   Zunrelated_tabsrD   r   Zrelated_tabsr   tabr   Z	page_elemeditorZedit_searchZdropdown_sortZbutton_sortorderZbutton_refreshZcheckbox_show_dirsr  r   r  r   r  ]  sD    
"*($$r  c                 C   s"   t j| } t j|}|| S r$   )r,   r-   rn   ro   )parent_path
child_pathr   r   r   path_is_parent  s    r  c                    sJ    fdd} j j|d j| jgg  jd  jD ]}|| q6d S )Nc           	         s   t |dkr(td fddjD S t| } | dk r<dn| } | t |krXt |d n| } || dkrj| nd }t|}t|\}}d}jD ]&}t fdd| D rd	} qq|sJ d
  dt||  fddjD S )Nr   z2There is no image in gallery to save as a preview.c                    s   g | ]}|  jqS r   r  r   r   r
  r   r   r     r   z2setup_ui.<locals>.save_preview.<locals>.<listcomp>r(   Fc                 3   s   | ]}t | V  qd S r$   )r  r   r   r   r   r'     r   z1setup_ui.<locals>.save_preview.<locals>.<genexpr>Tzwriting to z is not allowedc                    s   g | ]}|  jqS r   r  r  r
  r   r   r     r   )	rp   printr   r   r
   r   r/   r   r   )	r   imagesr    img_infoimagegeninforG   Z
is_allowedZ
extra_pager
  r   r   save_preview  s"    
zsetup_ui.<locals>.save_previewz:function(x, y, z){return [selected_gallery_index(), y, z]}r   )r   r  r   r   r  setup_ui)r   galleryr$  r  r   r
  r   r%    s    
r%  )r   )r   r   )r   r   r   )+os.pathr,   urllib.parserh   pathlibr   modulesr   r   r   r   modules.imagesr   r   gradior   rA   rE   fastapi.exceptionsr	   'modules.generation_parameters_copypaster
   modules.ui_componentsr   r   r   r   r   strr4   rC   rM   rR   rW   rX   r   r   r   r   Blocksr  r  r%  r   r   r   r   <module>   s6    Y	=