1、關鍵python依賴
(1)xformers:優化加速方案。它可以對模型進行適當的優化來加速圖片生成并降低顯存占用。缺點是輸出圖像不穩定,有可能比不開Xformers略差。
(2)GFPGAN:它是騰訊開源的人臉修復算法,利用預先訓練號的面部GAN(如styleGAN2)中封裝的豐富多樣的先驗因素進行盲臉(blind face)修復,旨在開發用于現實世界人臉修復的實用算法。
(3)CLIP:Contrastive Language-Image Pre-Training,多模態方向的算法。可以訓練出一個可以處理圖像和文本的模型,從而使得模型可以同時理解圖像和對圖像的描述。
(4)OPEN-CLIP:一個開源的clip實現。
(5)Pyngrok:Ngrok工具的python實現,可以實現內網穿透
2、核心目錄文件
(1)sd根目錄下的repositories
存放算法源碼
1)stable-diffusion-stability-ai:sd算法
2)taming-transformers:高分辨率圖像合成算法
3)k-diffusion:擴散算法
4)CodeFormer:圖片高清修復算法
5)BLIP:多模態算法
(2)sd根目錄/models
存放模型文件
3、Gradio使用說明
【stable diffusion webui源碼解析】-界面篇ui.py - 知乎
sd是基于gradio構建的,它是python庫,僅需幾行代碼就可以構造一個html界面。
測試例子:
gr.Interface是只有左右分列的布局,它有3個輸入參數:
參數1:處理函數,根據inputs中傳入的組件按照順序對應到函數的入參
參數2:組件信息
參數3:輸出的數據類型
4、webui之模型處理流程
(1)cleanup_models函數move模型文件
將models目錄下的文件移到相關子目錄下,比如ckpt文件和safetensors文件放到Stable-diffusion子目錄下。
(2)啟動SD模型setup_model流程
該模型位于:/data/work/xiehao/stable-diffusion-webui/models/Stable-diffusion
主要是通過list_models函數遍歷所有的模型的信息并存到checkpoint_alisases中。
第1步,查看sd/models/Stable-diffusion下是否有cpkt和safetensors結尾的文件,有則放入model_list列表中,沒有則從hugginface下載模型。
第2步,通過CheckpointInfo函數檢查model_list中每個模型的checkpoint信息。如果是safetensors文件,通過read_metadata_from_safetensors讀取文件信息。Safetensors模型的參數都存放在json中,把鍵值對讀出來存放到metadata字段中。
第3步,最后把每個模型根據{id : 模型對象}的鍵值對存放到checkpoint_alisases全局變量中。
(3)啟動codeformer模型的setup_model流程
該模型位于:/data/work/xiehao/stable-diffusion-webui/models/Codeformer
主要將Codeformer初始化之后的實例放到shared.face_restorers列表中。在此過程中并沒有將模型參數裝載到Codeformer網絡中。
(4)啟動GFPGAN模型的setup_model流程
(5)遍歷并加載內置的upscaler算法
這些算法位于:/data/work/xiehao/stable-diffusion-webui/modules
遍歷該目錄下_model.py結尾的文件,通過importlib.import_module()進行加載,這一步未看到實際作用。
初始化以下放大算法[<class 'modules.upscaler.UpscalerNone'>, <class 'modules.upscaler.UpscalerLanczos'>, <class 'modules.upscaler.UpscalerNearest'>, <class 'modules.esrgan_model.UpscalerESRGAN'>, <class 'modules.realesrgan_model.UpscalerRealESRGAN'>],其中第1個沒任何算法,第2-4是img.resize()方法實現的,第5、6個需要單獨加載模型,數據都以UpscalerData格式存放,其中該對象的local_data_path存放了模型的本地地址信息。
比如:shared.sd_upscalers[5].local_data_path為:
'/data/work/xiehao/stable-diffusion-webui/models/RealESRGAN/RealESRGAN_x4plus_anime_6B.pth'
(6)加載py執行腳本load_scripts
遍歷sd根目錄/scripts下的py腳本 以及 extensions下各擴展組件的py腳本,放到scripts_list變量中,格式如下:ScriptFile(basedir='/data/work/xiehao/stable-diffusion-webui/extensions/sd-webui-controlnet', filename='processor.py', path='/data/work/xiehao/stable-diffusion-webui/extensions/sd-webui-controlnet/scripts/processor.py')
遍歷并導入scripts_list中的類型為Script或ScriptPostprocessing的py文件:
Load_module(path)加載第三方組件時可能會輸出日志信息:
?
(7)遍歷VAE模型
目前沒有裝任何vae模型
(8)加載模型load_model
Select_checkpoint()函數,獲取sd模型信息,majicmixRealistic_v4.safetensors/majicmixRealistic_v4.safetensors [d819c8be6b]
do_inpainting_hijack函數。設置PLMSSampler的p_sample_plms。關于該方法,重建圖片的反向去噪過程的每一步的圖片都應用了該方法。
get_checkpoint_state_dict函數。如果是safetensors則使用safetensors.torch.load_file加載模型參數,否則使用torch.load加載模型參數。加載到pl_sd的dict類型變量中。
?? pl_sd字典做進一步處理:如果最外層是state_dict的key,則取該key下的value。此時pl_sd下就是模型各個節點名及對應的weights值。然后替換下面的key值:
?????? find_checkpoint_config函數。先從模型目錄下找下yaml配置文件,如果沒有則執行guess_model_config_from_state_dict函數,即從模型參數中獲取模型配置,最后返回/data/work/xiehao/stable-diffusion-webui/configs/v1-inference.yaml作為配置文件,信息如下:
?????? 接著用OmegaConf.load加載yaml文件,然后通過/data/work/xiehao/stable-diffusion-webui/repositories/stable-diffusion-stability-ai/ldm/util.py(82)instantiate_from_config()加載yaml信息獲得model。具體步驟為:
步驟1,通過yaml的target信息,可以知道model為ldm.models.diffusion.ddpm的LatentDiffusion類。模型的源碼位于:sd根目錄/modules/models/diffusion/ddpm_edit.py。
?????? 步驟2,通過getattr(module的obj,class_name)獲取model的類。
load_model_weights函數,將模型參數加載到模型中。通過model.load_state_dict(state_dict, strict=False)加載。因為程序參數no_half為false,所以模型量化需要從float32變為半精度tensor,half()的時候不對vae模塊做處理。Vae模塊為model.first_stage_model部分,所以先存到一個臨時變量,half()量化完成后再賦值回去。Vae最后再單獨變為float16。然后把模型放到cuda上。
Hijack函數,處理用戶輸入的embedding信息。假如給一個初始值,通過SD會生成未知的東西,我們通過添加額外的信息(比如prompts)讓sd朝著我們想要的方向生成東西,這個就是劫持的功能,劫持是在embeddings層的。模型的embedding類為:transformers.models.clip.modeling_clip.CLIPTextEmbeddings,它的token_embeddings類為:torch.nn.modules.sparse.Embedding。
針對prompts的embedding處理類為:FrozenCLIPEmbedderWithCustomWords。約有4.9W個token。然后針對token的權重進行處理,普通單詞為1.0, 中括號則除以1.1,小括號則乘以1.1.
指定優化方法apply_optimizations,通過xformers工具優化sd模型中的CrossAttention。(跨注意力機制是一種擴展自注意力機制的技術。自注意力機制是一種通過計算查詢query、鍵key和值value之間的關聯度來為輸入序列中的每個元素分配權重的方法,而跨注意力機制則通過引入額外的輸入序列來融合兩個不同來源的信息以實現更準確的建模)。
load_textual_inversion_embeddings函數,加載根目錄/embeddings下的embedding文件。加載[('/data/work/xiehao/stable-diffusion-webui/embeddings', <modules.textual_inversion.textual_inversion.DirWithTextualInversionEmbeddings object at 0x7ff2900b39d0>)]兩個下的embeddings信息。比如:badhandv4、easynegative、EasyNegativeV2、ng_deepnegative_v1_75t等。
model_loaded_callback函數,遍歷callback_map['callbacks_model_loaded']所有的回調函數,然后把sd_model模型傳進去依次執行這些回調函數。比如/data/work/xiehao/stable-diffusion-webui/extensions/a1111-sd-webui-tagcomplete/scripts/tag_autocomplete_helper.py的get_embeddings方法,/data/work/xiehao/stable-diffusion-webui/extensions-builtin/Lora/scripts/lora_script.py的assign_lora_names_to_compvis_modules方法。
5、頁面布局
基于gradio編寫,界面入口函數為modules/ui.py的create_ui()。
未完待續