Modal & Diffusers の最強タッグで画像生成する : GPUが無いPCで画像生成する方法

初めに

最近巷では画像生成AIというものがかなりHotであり、様々なアプリやサイトができているが、その殆どがサーバー上で処理した結果を送り返している。興味や知識はある人は自分のPCで画像生成をしたりしているが、GPU、さらにはVRAMまでいる為、実行できる人は少ない。
今回は、低性能なPCやGPUを積んでいないPCにおいて画像生成をする方法を、技術と一緒に書いていきます。

Modal について

公式サイトによると、

Your end-to-end stack
for cloud compute
Modal lets you run or deploy machine learning models, massively parallel compute jobs, task queues, web apps, and much more, without your own infrastructure.

つまりは、GPU等の計算資源を貸してくれるサービスです。
今回はこれを使って、画像生成プログラムを書いていきます。

始め方

まず上の公式サイトでGithub アカウントからログインします。
すると認証待ちになるので、待ちます。
通常の場合１日程度で認証されるのでそこからです。気長に待ちましょう。
また、計算時にはその計算量によって課金されますが、最初に$30.00 配られるため、クレジットカード等を登録する必要はありません。
それ以上使う場合は登録してください

にアクセスできるようになったら、

pip install modal-client
modal token new

を実行しましょう。modal token new でブラウザが開きます。
自分のPCのmodal-clientにトークンが保存されて使用できる状態になります。

使用する準備が整いました

Diffuser とは

Diffuserは簡単に言えば、画像生成を簡単に書けるようにしてくれたmoduleです。
HuggingFaceが開発しているやつです。

画像生成する方法

早速作業していきましょう。
楽しいコーディングの時間です！

0. 概要

ModalでDiffuserを使って画像生成する場合のリファレンスを書きます。
そのため、Pythonを理解していない人や、画像生成だけしたい人は的外れの記事かなと思います。すみません。

今回することをざっくり説明します。
自分のPC上で Modal-Client を立ち上げ、リモート上で計算を行う関数内でDiffuserを使って画像を生成します。
DiffuserはHuggingFaceから直接モデルをインポートできるみたいなのでそれを利用します。

1. HuggingFace Access Token の作成

このトークンを使ってHuggingFace.coから画像生成の学習済みDiffuserモデルをインポートしてきます。
通常自分のPCで実行する場合はそのモデルをダウンロードしますが、このトークンを使うとその必要がなくなります。
学習済みモデルは数十ギガを超えるため、ダウンロードは困難ですが、これを使うことによって大分処理が簡単になります。

以下の手順により、HuggingFace Access Token を作成してください。

1. HuggingFace にアクセスして左側の Settings を開きます
2. これまた左の Access Token を選択して、New token から新しいトークンを作成してください。
3. 名前はたんでもいいですがわかりやすいものを入れてください。今回であれば DiffuserInModal とかに設定しました。
4. Role は読み取りだけなので、read のままで構いません。

作成できたらトークンをコピーしておきます。

2. Modal でシークレットの作成

シークレット、要は実行時の環境変数や、スナップショット等を保存するために必要なものです。
無くても実行はできますが、pip のインストールを保持してくれるため作っておいて損はありません。

以下の手順で作成します。

1. Modal Home の上のタブから SECRETS を選択します。
2. NEW SECRET を選択し、一番下のCustom を選びます。
3. 環境変数の設定です。ここに先ほどコピーした HuggingFaceのトークンを保存しておきます。
4. KEY は名前なのでわかりやすいように HUGGINGFACE_TOKEN としておきましょう
5. VALUE は先ほどコピーしたトークンを張り付けます。
6. 下の NEXT を押します。
7. このシークレットの名前を入力します。なんでもいいですが、今回は my-custom-secret のままで行きます。
8. CREATE を押し、FINISH を押します。

これでシークレットが作成できました。

3. コーディング

準備ができました。
[1. CODE] で実行できます。

modal run [File Name]

のように実行します。

1. CODE

import os
import io
import modal

stub = modal.Stub()

@stub.function(
    image=modal.Image.debian_slim().pip_install("torch", "diffusers[torch]", "transformers", "ftfy", "discord"),
    secret=modal.Secret.from_name("my-custom-secret"),
    gpu="a10g",
    timeout=6000)
def run_diffusion(prompt: str, negative_prompt: str, height: int, width: int, steps: int, scale: float):
    import torch
    from diffusers import StableDiffusionPipeline

    defautPrompt = [
        'masterpiece, best quality, ',
        'lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, '
    ]

    pipe = StableDiffusionPipeline.from_pretrained(
        "Linaqruf/anything-v3.0",
        use_auth_token=os.environ["HUGGINGFACE_TOKEN"],
        torch_dtype=torch.float32,
    ).to("cuda")

    image = pipe(
        prompt=defautPrompt[0] + prompt,
        height=height,
        width=width,
        num_inference_steps=steps,
        guidance_scale=scale,
        negative_prompt=defautPrompt[1] + negative_prompt,
        # num_images_per_prompt=4,
        # eta=0.0,
        # generator=None,
        # latents=None,
        # prompt_embeds=None,
        # negative_prompt_embeds=None,
        # output_type="pil",
        # return_dict=True,
        # callback=None,
        # callback_steps=1,
    ).images[0]

    buf = io.BytesIO()
    image.save(buf, format="PNG")
    img_bytes = buf.getvalue()

    return img_bytes

@stub.local_entrypoint
def main():
    img_bytes = run_diffusion.call(
        prompt="",
        negative_prompt="",
        height=768, width=512, steps=28, scale=12.0)
    with open("output.png", "wb") as f:
        f.write(img_bytes)

上記コードはModalでDiffuserを動作させ、画像のPNGデータを返し、PCに保存するコードです。

2. 解説

1. Modal コード

import os
import io
import modal

stub = modal.Stub()

@stub.function(
    image=modal.Image.debian_slim().pip_install("torch", "diffusers[torch]", "transformers", "ftfy", "discord"),
    secret=modal.Secret.from_name("my-custom-secret"),
    gpu="a10g",
    timeout=6000)
def run_diffusion(prompt: str, negative_prompt: str, height: int, width: int, steps: int, scale: float):
    pass

@stub.local_entrypoint
def main():
    run_diffusion.call(
        prompt="",
        negative_prompt="",
        height=768, width=512, steps=28, scale=12.0)

1. main 関数からロードされます。
2. run_diffuser 関数はModalのリモート上で計算を行うための関数です。ここに処理を書くことで、Modalが計算して返します。
3. この関数を呼ぶには、.call(args...) 関数を使用して呼び出す必要があります。
4. @stub.function() はリモート側の設定です。ここに使用するためのモジュールや、シークレット名、使用するGPU等を書きます。

2. Diffuser コード

    import torch
    from diffusers import StableDiffusionPipeline

    defautPrompt = [
        'masterpiece, best quality, ',
        'lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, '
    ]

    pipe = StableDiffusionPipeline.from_pretrained(
        "Linaqruf/anything-v3.0",
        use_auth_token=os.environ["HUGGINGFACE_TOKEN"],
        torch_dtype=torch.float32,
    ).to("cuda")

    image = pipe(
        prompt=defautPrompt[0] + prompt,
        height=height,
        width=width,
        num_inference_steps=steps,
        guidance_scale=scale,
        negative_prompt=defautPrompt[1] + negative_prompt,
        # num_images_per_prompt=4,
        # eta=0.0,
        # generator=None,
        # latents=None,
        # prompt_embeds=None,
        # negative_prompt_embeds=None,
        # output_type="pil",
        # return_dict=True,
        # callback=None,
        # callback_steps=1,
    ).images[0]

    buf = io.BytesIO()
    image.save(buf, format="PNG")
    img_bytes = buf.getvalue()

1. AIの画像生成をするためのコードです。

Reference