引子
博主很長一段時間都沒有發文，確實是在忙一些技術研究。
如標題所示，本篇博文主要把近段時間的研究工作做一個review。
看過各種相關技術的公關文章，林林總總，水分很多。
也確實沒有多少人能把一些技術細節用一些比較通俗的語言闡述清楚。
故此，再一次冠以大話為題，對AI繪畫主要是stable diffusion做一個技術梳理。
如何學習以及相關資源
相信很多朋友都想入門到這個技術領域搗騰搗騰，
而擺在眼前的確是一條臭水溝。
為什麼這麼說，讓我們來看一些資料。
-
Hardware: 32 x 8 x A100 GPUs
-
Optimizer: AdamW
-
Gradient Accumulations: 2
-
Batch: 32 x 8 x 2 x 4 = 2048
-
Learning rate: warmup to 0.0001 for 10,000 steps and then kept constant
Hardware Type: A100 PCIe 40GB
-
Hours used: 150000
-
Cloud Provider: AWS
-
Compute Region: US-east
-
Carbon Emitted (Power consumption x Time x Carbon produced based on location of power grid): 11250 kg CO2 eq.
摘自：CompVis/stable-diffusion-v1-4 · Hugging Face
該模型是在亞馬遜雲端計算服務上使用256個NVIDIA A100 GPU訓練，共花費15萬個GPU小時，成本為60萬美元
摘自：Stable Diffusion - 維基百科，自由的百科全書 (wikipedia.org)
這個資料就是一個勸退警告，但是由於效果太過於“嚇人”，所以飛蛾撲火，全世界都打起架來了。
當然，剛開始學習，就直接奔著最終章去，確實也不是很現實。
隨著這個領域的爆火，各種資源爆炸式增長。
以下是博主給出的一部分參考資源，便於參閱。
相關整合資源:
第三方:
Generative Deep Learning (keras.io)
AUTOMATIC1111/stable-diffusion-webui: Stable Diffusion web UI (github.com)
官方:
CompVis/stable-diffusion: A latent text-to-image diffusion model (github.com)
想要開箱式快速上手，建議把keras社群的這個生成花朵的玩具跑起來感受一下。
Denoising Diffusion Implicit Models (keras.io)
這個是一個非常簡潔的實現，麻雀雖小五臟俱全，可以快速熱身起來。
Stable Diffusion的基本原理
這裡博主並不打算展開講解過多，
只做一個大話概覽闡述，便於快速入門瞭解。
若還有疑惑，大家可以參閱其他的資源。
瞭解原理，僅僅閱讀程式碼肯定是遠遠不夠的，
但是不閱讀程式碼，你就不知道具體的實現細節。
官方實現，那個程式碼倉庫真是一座屎山，亂七八糟的。
所以極佳的閱讀版本是keras社群的實現。
keras-cv/keras_cv/models/stable_diffusion
stable_diffusion的主要元件如下:
1.文案編碼器：
負責對輸入的文字進行特徵編碼，用於引導 diffusion 模型進行內容生成。
2.潛在特徵編碼器：
將圖片編碼成潛在特徵
3.潛在特徵解碼器：
將潛在特徵解碼成圖片
4.diffusion 模型:
由噪聲和文字語義生成目標潛在特徵。
用通俗的話來說編碼器就是壓縮，解碼器就是解壓，diffusion 模型就是編輯壓縮的資訊，而文案是引導編輯的方向。
5.輔助元件:
針對文案輸入的編碼轉換以及預處理。
簡而言之就是將文字轉換成數字。
噪聲規劃，主要用於訓練和合成的加噪和去噪的比率計算。
整體的架構串起來就是stable_diffusion
整個使用的流程大概是這個樣子的：
文字生成圖片:
文字 -> clip_tokenizer -> text_encoder -> diffusion_model -> decoder
圖片生成圖片:
[[圖片 -> image_encoder] + [文字 -> clip_tokenizer -> text_encoder]] -> diffusion_model -> decoder
這是主流的兩種做法，還有透過mask進行內容修復，以及透過其他模組輔助生成的，例如生成小姐姐之類等等。
列完元件，
你就會發現涉及的技術並不簡單。
說好的，要講原理的。
一句話描述： 這是一個訊號壓縮和解壓的演算法。
其中的訊號是透過噪聲擴散建的模，建模後符合正向和逆向的數學規律。
而在正向和逆向中都可以注入噪聲或者改變噪聲來達到訊號引導重建。
臨時開個小叉：
我們知道:
聲音檔案wav進行餘弦變換就可以壓縮成Mp3
點陣圖檔案bmp進行餘弦變換就可以壓縮成jpg
那麼你說如果我們對餘弦變換進行數學編輯，是不是可以直接編輯目標的mp3或者目標的jpg
你答對了，是可以的。
最經典的做法就是採用餘弦變換對訊號進行濾波，然後可以做到降噪，也就可以做到音訊降噪或者圖片降噪。
我們回到主題上來，基於噪聲建模之後，有什麼好處，這樣做了之後故事變得更加有意思的。
只要編輯後的資料符合噪聲分佈，那理論上來說，我們可以採用正向或者逆向擴散的方式對任意時刻的訊號進行編輯。
而Stable Diffusion 就是這樣一個技術方案。
訓練時候它將資料透過noise_scheduler進行diffusion_model正向擴散分治壓縮，
生成使用的時候透過diffusion_model進行逆向擴散編輯解壓，
最終達到可以生成任意訊號組合的圖片。
我知道你有疑問了，那是不是在生成的時候可以在任意時候插入資訊，或者是正向擴散和逆向擴散混合著來。
沒錯，完全可以的。
圖片生成圖片，就是在逆向擴散過程中插入一個我們預設的圖片資訊節點與噪聲進行混合，然後透過文字語義引導合成的過程。
知道了技術原理，我相信你們跟我一樣會有一個非常大膽的假設。
假設把全世界的電腦連線起來，然後採用擴散分治，那完全可以讓世界上任何一個地區任何一個人的一臺機器，幫你把某個節點的擴散編輯給算了。
那這個事情就恐怖了，全世界都是p2p 算力共享了，人工智慧的大時代就到來了。
就問一個問題，現在上船還來得及嗎？
博主也不知道，只是這個事已經是鐵板釘釘了，趨勢不可逆。
關鍵演算法以及相關最佳化
前面講到技術原理，但是到底其中關鍵演算法是什麼？
答案是：跨模態注意力。
一切都由谷歌的一篇論文開始。
[1706.03762] Attention Is All You Need (arxiv.org)
在這篇論文出現後，百花齊放，各種attention, 各種transformer層出不窮。
transformer架構的提出，給出不同維度資料建模的可行性。
關於這方面的發展和資源，博主只給出如下資源。
一個值得關注的技術github:
lucidrains (Phil Wang) · GitHub
lucidrains是一個非常勤奮且高產的人，幾乎所有第三方的transformer實現都有他的影子，當然也包括stable_diffusion。
在stable_diffusion中也是因為使用transformer所以帶來了嚴重的計算資源的消耗。
針對stable_diffusion 的演算法最佳化，各大廠也是勤奮的。
附相關資訊:
英特爾
Accelerating Stable Diffusion Inference on Intel CPUs (huggingface.co)
高通
World’s first on-device demonstration of Stable Diffusion on an Android phone | Qualcomm
蘋果
Stable Diffusion with Core ML on Apple Silicon - Apple Machine Learning Research
谷歌
當然我們這裡講的最佳化指的是使用階段的最佳化，並不是訓練階段的最佳化。
這是完全不同的兩件事，而現在國內泛指的stable_diffusion訓練也不是指的從零訓練，
而是在stable_diffusion的基礎上進行二次訓練微調。
畢竟訓練成本擺在那裡，沒有幾家公司會這麼豪氣去復現，包括國內某些大廠。
訓練最佳化展開一時半會也講不完，所以博主著重講一下，使用階段的最佳化。
參考資訊來自:
Optimizations · AUTOMATIC1111/stable-diffusion-webui Wiki (github.com)
| commandline argument | explanation |
|---|---|
--opt-sdp-attention |
Faster speeds than using xformers, only available for user who manually install torch 2.0 to their venv. (non-deterministic) |
--opt-sdp-no-mem-attention |
Faster speeds than using xformers, only available for user who manually install torch 2.0 to their venv. (deterministic, slight slower than --opt-sdp-attention) |
--xformers |
Use xformers library. Great improvement to memory consumption and speed. Will only be enabled on small subset of configuration because that's what we have binaries for. Documentation |
--force-enable-xformers |
Enables xformers above regardless of whether the program thinks you can run it or not. Do not report bugs you get running this. |
--opt-split-attention |
Cross attention layer optimization significantly reducing memory use for almost no cost (some report improved performance with it). Black magic. On by default for torch.cuda, which includes both NVidia and AMD cards. |
--disable-opt-split-attention |
Disables the optimization above. |
--opt-sub-quad-attention |
Sub-quadratic attention, a memory efficient Cross Attention layer optimization that can significantly reduce required memory, sometimes at a slight performance cost. Recommended if getting poor performance or failed generations with a hardware/software configuration that xformers doesn't work for. On macOS, this will also allow for generation of larger images. |
--opt-split-attention-v1 |
Uses an older version of the optimization above that is not as memory hungry (it will use less VRAM, but will be more limiting in the maximum size of pictures you can make). |
--medvram |
Makes the Stable Diffusion model consume less VRAM by splitting it into three parts - cond (for transforming text into numerical representation), first_stage (for converting a picture into latent space and back), and unet (for actual denoising of latent space) and making it so that only one is in VRAM at all times, sending others to CPU RAM. Lowers performance, but only by a bit - except if live previews are enabled. |
--lowvram |
An even more thorough optimization of the above, splitting unet into many modules, and only one module is kept in VRAM. Devastating for performance. |
*do-not-batch-cond-uncond |
Prevents batching of positive and negative prompts during sampling, which essentially lets you run at 0.5 batch size, saving a lot of memory. Decreases performance. Not a command line option, but an optimization implicitly enabled by using --medvram or --lowvram. |
--always-batch-cond-uncond |
Disables the optimization above. Only makes sense together with --medvram or --lowvram |
--opt-channelslast |
Changes torch memory type for stable diffusion to channels last. Effects not closely studied. |
--upcast-sampling |
For Nvidia and AMD cards normally forced to run with --no-half, should improve generation speed. |
從這個表單上，我們能得出的資訊就是最佳化的方向，基本都是視訊記憶體佔用。
除了視訊記憶體佔用外的最佳化也有幾條路子可以走。
1.最佳化運算元，平行計算提速
例如：
將sigmoid 替換成等價的tanh，因為數學計算中tanh的計算更加簡單
2.量化精度，採用fp16或更低的精度求近似解
fp16帶來的一些計算或溢位取值範圍，所以如果採用量化精度的方案，必須考慮計算的值域問題。
3.微調蒸餾，將模型中的某些耗時計算蒸餾
例如將模型中的transformer採用等價的近似實現替換掉或蒸餾一下。
4.最佳化擴散取樣演算法，基於擴散的數學先驗，加速求解
擴散取樣需要多步取樣才能達到理想的視覺效果，改進取樣演算法，減少取樣步數，是個非常理想的最佳化方向，但是並沒有那麼容易做到。
而上表中的xformers就是一個針對了transformer進行記憶體最佳化和計算最佳化的開源方案。
對效能最佳化感興趣的朋友，建議好好閱讀一下xformers的實現。
視訊記憶體方面的最佳化主要研讀以下兩篇論文：
[2112.05682v3] Self-attention Does Not Need O(n^2) Memory
很多聲稱對stable_diffusion做了最佳化的廠家，基本上都是在模型轉換或者編譯階段將這兩篇論文的思路做了應用而已，沒有各公眾號文章說得那麼玄乎。
另外說一句，公眾號文章很多水分和誇大，看看笑笑就好，不用較真。
博主也在TensorFlow下做了相關的演算法移植工作，實現起來並不麻煩。
當然除了以上提到最佳化方案之外還有不少思路，只是其中有一些通用性不強或者說存在特定的侷限性。
最佳化這件事，見仁見智。
落地部署的林林總總
待有空再講講部署使用模型的一些道道。
主要涉及模型轉換，精度處理，問題排查。
stable_diffusion 部署和使用，
展開說也確實三言兩語說不完。
博主成功將stable_diffusion移植到低端手機跑起來，並且成功出圖。
不過還有很多工作要做才能達到在普通手機上秒級出圖，路漫漫其修遠兮。
這方面博主還在努力，歡迎有相關經驗的朋友一起探討。
發展和研究的建議
從0訓練復現stable_diffusion的可行性，是完全可能的，
只不過訓練成本不會低的，要做的工作也不少，
博主，已經實現了一套全新的訓練方案，改進了訓練中的各個環節，包括最佳化器演算法，都是獨創的，但是奈何算力資源仍然不夠用，若有機構願意贊助，歡迎來郵洽談。
stable_diffusion原本的架構非常的冗餘，並且很野蠻，所謂大力出奇跡。
除了之前詬病說的生成人手問題等，還有很大的改進空間，
希望有更多的朋友加入進來，一起擁抱趨勢，擁抱未來。
謹以此文，與同道人士，共勉之。
若有其他相關問題或者相關技術需求歡迎來郵聯絡。
郵箱地址是: gaozhihan@vip.qq.com