📎

[WebUI] Stable DiffusionベースモデルのCLIPの重みを良いやつに変更する

2022/11/29に公開約9,100字3件のコメント

簡単な記事なので前置きは省略します。

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を置き換えるスクリプトを有志の方が作成されたようなので、こちらを使用してみてください。

https://github.com/bbc-mc/sdweb-clip-changer

以下の文章は古いやり方の紹介と、その検証結果です。

.
.
.

本題

Stable Diffusion WebUIのフォルダを開き、modules/sd_hijack.pyの83-90行目を探してください。

https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/0b5dcb3d7ce397ad38312dbfc70febe7bb42dcc3/modules/sd_hijack.py#L83-L90

そして、このように変更します。

    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(...)の中のTrueFalseに書き換えてください。挙動が元に戻り、モデルが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でのテキスト条件付けにおいても同じことが言えるかは分かりませんが、とりあえず試してみたのでその方法を追記します。

まず、下のコードを変更します。

変更前

https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/0b5dcb3d7ce397ad38312dbfc70febe7bb42dcc3/modules/sd_hijack.py#L83-L90

変更後

    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変更の影響について、申し訳ありませんが自分の知識では断言できません…。
推測の範囲で述べさせていただきます。

  • HyperNetworks
    • 単純にCLIPを変更し、追加学習を行わない場合と同様の良い影響(プロンプトの効きが若干良くなる)が出るのではないかと思います。
  • Textual Inversion
    • Textual Inversionの埋め込みの質はおそらく変わりません。一方、変更前のCLIPで作成したTextual Inversionの重みがうまく使えなくなる可能性があります(未検証です)。
  • Dreambooth
    • Dreambooth後に出力される重みファイルの中に、変更したCLIPを再学習したものが含まれるようになるはずなので、何かしらの影響は出ると思います。それが良いものか悪いものかは分かりませんが、直感では良いものになると思っています。なお、Dreambooth後は必ずこちらを行っていただく必要があるので、ご注意ください。

      もし何か不具合がある場合、もしくはDreamboothを行ったモデルを使用する場合は、def hijack(...)の中のTrueをFalseに書き換えてください。挙動が元に戻り、モデルがckpt内にあるCLIPを使用するようになります。

ふわっとした回答になってしまって大変恐縮ですが、もし1ミリでも参考になれば幸いです。

詳しくお答え頂きありがとうございます!

ログインするとコメントできます