a
    d##                     @   s  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m	Z	m
Z
 d dlZz8d dlmZmZmZmZmZmZmZ d dlmZ dZW n ey   dZY n0 zd dlZdZW n ey   dZY n0 d	d
lmZmZmZ d	dlm Z  dZ!dZ"dZ#e$ee% dddZ&d:e e%e$e
e'e$f e'dddZ(d;ee% e$e$ee$ ee$ e'e'ee% e
e'e$f d	ddZ)d<e$e$e$ee	e*df  ee	e*df  e$ee$ ee$ e'e'ee% dddZ+e%e$ddd Z,e-d!kre j.d"d#Z/e/j0d$e$d%d& e/j0d'e$d(d& e/j0d)e$d*d& e/j0d+e$dd, e/j0d-e*d.dd/d0d1 e/j0d2e*d.dd3d4d1 e/1 Z2e3d5e2j4 d6e2j5 d7e2j6  e+e2j4e2j5e2j6e2j7e2j8e2j9d8 e3e2j4 d9 dS )=    N)Path)TemporaryDirectory)OptionalTupleUnion)create_repoget_hf_file_metadatahf_hub_download
hf_hub_urlrepo_type_and_id_from_hf_idupload_folderlist_repo_files)EntryNotFoundErrorTF   )create_model_from_pretrainedget_model_configget_tokenizer)HFTokenizerzopen_clip_pytorch_model.binzopen_clip_model.safetensorszopen_clip_config.json)config_pathmodel_configc                 C   sZ   | j j| j jd}||d}|d }tj||dd W d    n1 sL0    Y  d S )N)meanZstd)Z	model_cfgpreprocess_cfgw   )indent)visual
image_mean	image_stdopenjsondump)modelr   r   r   Z	hf_configf r#   a/var/www/html/stable-diffusion-webui/venv/lib/python3.9/site-packages/open_clip/push_to_hf_hub.pysave_config_for_hf'   s    r%   )	tokenizerr   save_directorysafe_serializationskip_weightsc           	      C   s   t }t|}|jddd |st|  }|du s6|dkrTtsBJ dtj||t  |du sd|dkrtt	||t
  || || }t| ||d d S )NT)exist_okparentsbothz-`pip install safetensors` to use .safetensorsF)r   )HF_CONFIG_NAMEr   mkdirZ
state_dict_has_safetensorssafetensorstorchZ	save_fileHF_SAFE_WEIGHTS_NAMEsaveHF_WEIGHTS_NAMEZsave_pretrainedr%   )	r!   r&   r   r'   r(   r)   Zconfig_filenameZtensorsr   r#   r#   r$   save_for_hf9   s    
r5   	Add model)	r   repo_idcommit_messagetokenrevisionprivate	create_pr
model_cardr(   c              
   C   sH  t |tstd}t|||dd}t|\}}}| d| }d}i }ztt|}d}W n. ty } ztd| W Y d }~n
d }~0 0 ztt	|d|d d}W n t
y   d}Y n0 t p}t| ||||
d	 |s|	pi }	|dd
 }t|d }t|	|}|| t|||||dW  d    S 1 s:0    Y  d S )Nzopenai/clip-vit-large-patch14T)r9   r;   r*   /FzRepo does not existz	README.md)r7   filenamer:   )r&   r   r'   r(   )r7   Zfolder_pathr:   r<   r8   )
isinstancer   r   r   setr   	Exceptionprintr   r
   r   r   r5   splitr   generate_readme
write_textr   )r!   r&   r   r7   r8   r9   r:   r;   r<   r=   r(   repo_url_Z
repo_owner	repo_nameZrepo_existsZ
repo_fileseZ
has_readmetmpdir
model_nameZreadme_pathreadme_textr#   r#   r$   push_to_hf_hubT   sL    
 
	

rO   fp32.)
pretrainedr7   	precisionr   r   r8   r9   r:   r;   r<   r=   c                 C   sP   t | ||||d\}}t| }|s&J t| }t||||||||	|
|dd d S )N)rQ   rR   r   r   r,   )r!   r&   r   r7   r8   r9   r:   r;   r<   r=   r(   )r   r   r   rO   )rM   rQ   r7   rR   r   r   r8   r9   r:   r;   r<   r=   r!   Zpreprocess_evalr   r&   r#   r#   r$   push_pretrained_to_hf_hub   s.    
rS   )r=   rM   c           	      C   s  d}|d7 }|d7 }|d7 }|d|  dd d7 }d	| v rld
| d	 v rl|d7 }|d| d	 d
   d7 }|d7 }|d| d7 }d| v r|d| d  d7 }d	| v rb|d7 }| d	  D ]\}}t|ttfr|d| d7 }|D ]}|d| d7 }qqt|trJ|d| d7 }| D ] \}}|d| d| d7 }q&q|d| d| d7 }qd| v r|d7 }|| d 7 }|d7 }d| v r|d7 }|| d 7 }|d7 }d| v r|d7 }t| d ttfs| d g}n| d }|D ]}|d| d7 }q|S )Nz---
ztags:
- clip
zlibrary_name: open_clip
z-pipeline_tag: zero-shot-image-classification
z	license: licenseZmit
detailsZDatasetz
datasets:
z- z# Model card for descriptionz
## Model Details
z- **z:**
z  - z: z:** usagez
## Model Usage
Z
comparisonz
## Model Comparison
Zcitationz
## Citation
z
```bibtex
z
```
)getloweritemsrA   listtupledict)	r=   rM   rN   kvviZkiZ	citationscr#   r#   r$   rF      sP    



rF   __main__zPush to Hugging Face Hub)rW   z--modelzName of the model to use.)typehelpz--pretrainedzHUse a pretrained CLIP model weights with the specified tag or file path.z	--repo-idz6Destination HF Hub repo-id ie 'organization/model_id'.z--precision)rd   defaultz--image-mean+ZMEANz,Override default image mean value of dataset)rd   nargsrf   metavarre   z--image-stdZSTDz2Override default image std deviation of of datasetzSaving model z with pretrained weights z to Hugging Face Hub at )rR   r   r   z saved.)FF)r6   NNFFNF)	rP   NNr6   NNFFN):argparser   ospathlibr   tempfiler   typingr   r   r   r1   Zhuggingface_hubr   r   r	   r
   r   r   r   Zhuggingface_hub.utilsr   Z_has_hf_hubImportErrorZsafetensors.torchr0   r/   factoryr   r   r   r&   r   r4   r2   r-   strr^   r%   boolr5   rO   floatrS   rF   __name__ArgumentParserparseradd_argument
parse_argsargsrD   r!   rQ   r7   rR   r   r   r#   r#   r#   r$   <module>   s   $	

  
        
I         *0


 	