[WebUI] Stable DiffusionベースモデルのCLIPの重みを良いやつに変更する
簡単な記事なので前置きは省略します。
CLIPについて
Stable Diffusionベースのモデルは、画像生成に際してテキストで条件づけを行う（a.k.a. txt2img）場合、CLIPという基盤モデルでテキストをベクトルに変換し、生成処理に渡しています。
つまり、入力したテキストをいかに的確に生成処理のモデルに伝えられるのかはCLIP（のようなテキストエンコーダー）の性能にかかっているわけですが、実はStable Diffusionに組み込まれているCLIPは、CLIPの中でも最強性能のものではありません。
適当に、画像のグラフに登場する点はそれぞれCLIPのバリアントで、上にあるモデルほど性能がいいと考えてください。
ここで、Stable Diffusion V1系に組み込まれているCLIPはL/14ですが、CLIPの中で最も性能の良いL/14@336pxに一歩遅れをとっていることが分かります。
For the ViT-L/14 we also pre-train at a higher 336
pixel resolution for one additional epoch to boost performance similar to FixRes (Touvron et al., 2019). We denote
this model as ViT-L/14@336px. - CLIP元論文より
つまり、L/14@336pxを使った方がStable Diffusionの性能上がるやん、ということで、この記事ではStable Diffusion WebUI上でCLIPの重みを一番いいやつに置き換える方法を紹介します。
[2023/01/20追記]
Stable Diffusion WebUIの更新により、この記事のやり方はout of dateになりました。
より良い方法でモデルのCLIPを置き換えるスクリプトを有志の方が作成されたようなので、こちらを使用してみてください。
以下の文章は古いやり方の紹介と、その検証結果です。
.
.
.
本題
Stable Diffusion WebUIのフォルダを開き、modules/sd_hijack.pyの83-90行目を探してください。
そして、このように変更します。
def hijack(self, m, use_improved_clip=True):
if type(m.cond_stage_model) == ldm.modules.encoders.modules.FrozenCLIPEmbedder:
if use_improved_clip:
from transformers import CLIPTextModel, CLIPTokenizer
device = "cuda" if torch.cuda.is_available() else "cpu"
m.cond_stage_model.transformer = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14-336").to(device)
m.cond_stage_model.tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14-336")
model_embeddings = m.cond_stage_model.transformer.text_model.embeddings
model_embeddings.token_embedding = EmbeddingsWithFixes(model_embeddings.token_embedding, self)
m.cond_stage_model = sd_hijack_clip.FrozenCLIPEmbedderWithCustomWords(m.cond_stage_model, self)
elif type(m.cond_stage_model) == ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder:
m.cond_stage_model.model.token_embedding = EmbeddingsWithFixes(m.cond_stage_model.model.token_embedding, self)
m.cond_stage_model = sd_hijack_open_clip.FrozenOpenCLIPEmbedderWithCustomWords(m.cond_stage_model, self)
終わりです。
この記事を書く直前まではだいぶ時間をかけて難しいことをやっていましたが、さっき、これだけで上手くいくことに気づきました。
この変更によって勝手に内部で使用するCLIPがL/14@336pxに切り替わります。使用感は変わりませんし、[]や()によるトークンの強調機能も問題なく使えます。
もし何か不具合がある場合、もしくはDreamboothを行ったモデルを使用する場合は、def hijack(...)の中のTrueをFalseに書き換えてください。挙動が元に戻り、モデルがckpt内にあるCLIPを使用するようになります。
比較
以下のように、なんか難しそうなプロンプトで生成します。
シード値のようなものは全て揃えています。
a digital painting of kawaii 1girl school student, focus_on_face, black hair, twintails, blue eyes, white shirt, yellow headband, evening, perfect symmetrical pretty face, bangs, extremely detailed, best quality, 8k uhd
L/14 (デフォルト)
悪くありませんが、白シャツを指定したのにセーラー服を着ています。
L/14@336px (置き換え後)
上手くいっています。
この比較ではあまり威力を感じないかもしれませんが、個人的に構図指定のプロンプトを追加したときの生成に対する効きが目に見えて良くなったと感じました。
なお、CLIP自体の限界があるため、良いやつに変更したとしても上手くいかない場合はまあまああります。
おわりに
割とフリーランチでStable Diffusionベースの画像生成モデルを改善できるので、ぜひ。
.
.
.
[2022-12-06 追記] BONUS: 2つのCLIPモデルの平均をとる
- 注意：試す価値はあんまりありません
その昔、Disco Diffusionという、今とは別の手法でテキストによる条件付けを行い、Diffusion Modelでtxt2imgをやるPlaygroundがありました。
そこでは複数のCLIPモデルをアンサンブルして使うことが当たり前になっていて、例えばCLIPのResnet-50とViT-L/14@336pxを組み合わせて使うと、それぞれ片方1つだけで使用するよりも生成結果が良くなることが知られていました。
Stable Diffusionでのテキスト条件付けにおいても同じことが言えるかは分かりませんが、とりあえず試してみたのでその方法を追記します。
まず、下のコードを変更します。
変更前
変更後
def hijack(self, m, use_improved_clip=True):
if type(m.cond_stage_model) == ldm.modules.encoders.modules.FrozenCLIPEmbedder:
if use_improved_clip:
from transformers import CLIPTextModel, CLIPTokenizer
device = "cuda" if torch.cuda.is_available() else "cpu"
m.cond_stage_model.tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14-336")
m.cond_stage_model.transformer = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14-336").to(device)
model_embeddings = m.cond_stage_model.transformer.text_model.embeddings
model_embeddings.token_embedding = EmbeddingsWithFixes(model_embeddings.token_embedding, self)
###################################################################################################################
# 変更点1
transformer_2 = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14").to(device)
model_embeddings_2 = transformer_2.text_model.embeddings
model_embeddings_2.token_embedding = EmbeddingsWithFixes(model_embeddings_2.token_embedding, self)
m.cond_stage_model = sd_hijack_clip.FrozenCLIPEmbedderWithCustomWords_2(m.cond_stage_model, self, transformer_2)
###################################################################################################################
else:
model_embeddings = m.cond_stage_model.transformer.text_model.embeddings
model_embeddings.token_embedding = EmbeddingsWithFixes(model_embeddings.token_embedding, self)
m.cond_stage_model = sd_hijack_clip.FrozenCLIPEmbedderWithCustomWords(m.cond_stage_model, self)
elif type(m.cond_stage_model) == ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder:
m.cond_stage_model.model.token_embedding = EmbeddingsWithFixes(m.cond_stage_model.model.token_embedding, self)
m.cond_stage_model = sd_hijack_open_clip.FrozenOpenCLIPEmbedderWithCustomWords(m.cond_stage_model, self)
そして次に、modules/sd_hijack_clip.pyを開き、以下のコードを追加してください。
変更後
class FrozenCLIPEmbedderWithCustomWords_2(FrozenCLIPEmbedderWithCustomWordsBase):
def __init__(self, wrapped, hijack, additional_transformer):
super().__init__(wrapped, hijack)
self.tokenizer = wrapped.tokenizer
self.comma_token = [v for k, v in self.tokenizer.get_vocab().items() if k == ',</w>'][0]
self.token_mults = {}
tokens_with_parens = [(k, v) for k, v in self.tokenizer.get_vocab().items() if '(' in k or ')' in k or '[' in k or ']' in k]
for text, ident in tokens_with_parens:
mult = 1.0
for c in text:
if c == '[':
mult /= 1.1
if c == ']':
mult *= 1.1
if c == '(':
mult *= 1.1
if c == ')':
mult /= 1.1
if mult != 1.0:
self.token_mults[ident] = mult
self.id_start = self.wrapped.tokenizer.bos_token_id
self.id_end = self.wrapped.tokenizer.eos_token_id
self.id_pad = self.id_end
self.additional_transformer = additional_transformer
def tokenize(self, texts):
tokenized = self.wrapped.tokenizer(texts, truncation=False, add_special_tokens=False)["input_ids"]
return tokenized
def encode_with_transformers(self, tokens):
outputs = self.wrapped.transformer(input_ids=tokens, output_hidden_states=-opts.CLIP_stop_at_last_layers)
outputs_2 = self.additional_transformer(input_ids=tokens, output_hidden_states=-opts.CLIP_stop_at_last_layers)
if opts.CLIP_stop_at_last_layers > 1:
z = outputs.hidden_states[-opts.CLIP_stop_at_last_layers]
z = self.wrapped.transformer.text_model.final_layer_norm(z)
z_2 = outputs_2.hidden_states[-opts.CLIP_stop_at_last_layers]
z_2 = self.additional_transformer.text_model.final_layer_norm(z_2)
z = (z + z_2) / 2
else:
z = outputs.last_hidden_state
return z
def encode_embedding_init_text(self, init_text, nvpt):
embedding_layer = self.wrapped.transformer.text_model.embeddings
ids = self.wrapped.tokenizer(init_text, max_length=nvpt, return_tensors="pt", add_special_tokens=False)["input_ids"]
embedded = embedding_layer.token_embedding.wrapped(ids.to(devices.device)).squeeze(0)
return embedded
これで、2つのCLIP(ViT-L/14 & ViT-L/14@336px)から出てきた埋め込み表現の平均値によってStable Diffusionの条件付けが行われるようになります。
そして、実験結果は以下。
追加実験・同一プロンプトでシードを揃えて6枚ずつ
うーん、まあまあというか、なんか退行してない…？
改善されてるかどうかよく分からない上、追加でVRAMが占有されるので、現時点であんまりオススメはできません。
ViT-L/14からViT-L/14@336pxへの変更による改善のほうが断然際立つ実験結果になっていますね…。
Discussion
有用な記事有難く拝見させてもらっています。
素人質問で恐縮なのですがCLIPを上記方法で変更した場合、DreamBoothなどでの追加学習に良い影響があったりするのでしょうか？
リョウヘイさん、こんにちは。
記事をご覧いただきありがとうございます！
Dreamboothなどの追加学習におけるCLIP変更の影響について、申し訳ありませんが自分の知識では断言できません…。
推測の範囲で述べさせていただきます。
ふわっとした回答になってしまって大変恐縮ですが、もし1ミリでも参考になれば幸いです。
詳しくお答え頂きありがとうございます！