LCM-LoRAを使って画像生成APIを作成する

はじめに

前にLCMを使った画像生成APIを作成しました。

https://zenn.dev/midra_lab/articles/962f1d578fede0

その後にLCM LoRAという新しいものができたので、これを使って再度APIを作っていきたいと思います。こちらを使うことでモデルを自由に選択できるのでかなり幅が広がったのではないかと感じています
https://huggingface.co/blog/lcm_lora

環境

• modal
• diffusers==0.24.0

準備

まずは、ベースになるモデルを決めます。
LCM-LoRAを使う以上以下で作成されたモデルでしか現時点では使用できません。

• latent-consistency/lcm-lora-sdxl : SDXL 1.0用ののLCM LoRA
• latent-consistency/lcm-lora-sdv1-5 : Stable Diffusion 1.5用のLCM LoRA
• latent-consistency/lcm-lora-ssd-1b : SSD-1B用のLCM LoRA

そのため、いくつかのモデルを探していきます。
使えそうなモデルは、以下あたりがありそうです

• civitaiでの検索結果

https://civitai.com/models/4468/counterfeit-v30

https://civitai.com/models/28169/cuteyukimixadorable-style?modelVersionId=237159

今回は「counterfeit-v30」を使用します

APIの実装

モデルのダウンロード

ベースモデルを一度Storageに持ってくる必要があるので、ダウンロードを行います。
モデルのサイズも大きいので、urllibを使ってダウンロードを行います

import urllib.request
url = 'https://huggingface.co/gsdf/Counterfeit-V3.0/resolve/main/Counterfeit-V3.0_fix_fp16.safetensors'
save_name = 'Counterfeit-V3.0_fix_fp16.safetensors'
urllib.request.urlretrieve(url, save_name)

モデルのロード

ベースモデルとLCM-LoRAをロードします。今回は、i2iをしたいので、StableDiffusionImg2ImgPipeline.from_single_file でモデルをロードします

# パイプラインの準備
pipe = StableDiffusionImg2ImgPipeline.from_single_file(
"Counterfeit-V3.0_fix_fp16.safetensors",
variant="fp16",
torch_dtype=torch.float16
).to("cuda")

# LoRAウェイトの準備
pipe.load_lora_weights("latent-consistency/lcm-lora-sdv1-5")

# スケジューラの準備
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
    

画像の生成

画像を生成をするときに、LCM-LoRAの場合は、guidance_scale を0にする必要があるのでこちらを設定します

image = pipe(
prompt=prompt,
image=load_image(image_url),
num_inference_steps=4,
guidance_scale=0,  # guide_scaleの無効化
generator=torch.Generator("cpu").manual_seed(random_seed),
).images[0]
    

base64に変換

APIとして使用したいので、扱いやすいbase64に変換します。

import io
import base64

# バイトIOオブジェクトを初期化
buf = io.BytesIO()
# 画像をPNG形式でバイトIOオブジェクトに保存
image.save(buf, format="PNG")
# バイトIOオブジェクトをBase64エンコードする
return base64.b64encode(buf.getvalue()).decode('utf-8')

クライアント側

modalにDeployした際にURLが発行されるので、getでたたくだけになります。
pythonの場合は、GETは以下のように叩けます

import requests
response = requests.get(url)

コード全体

以下がAPIのコードになります。

from modal import Stub, web_endpoint
import modal

stub = Stub("generate_i2i_lora")


# 画像から画像を生成するAPI
@stub.function(
    image=modal.Image.debian_slim().pip_install("omegaconf", "diffusers==0.24.0", "torch", "transformers"),
    secret=modal.Secret.from_name("HUGGINGFACE_TOKEN"),
    gpu="t4",
    timeout=1000)
@web_endpoint()
def run():
    import urllib.request
    url = 'https://huggingface.co/gsdf/Counterfeit-V3.0/resolve/main/Counterfeit-V3.0_fix_fp16.safetensors'
    save_name = 'Counterfeit-V3.0_fix_fp16.safetensors'
    urllib.request.urlretrieve(url, save_name)

    from diffusers import StableDiffusionImg2ImgPipeline, LCMScheduler
    import torch

    # パイプラインの準備
    pipe = StableDiffusionImg2ImgPipeline.from_single_file(
        "Counterfeit-V3.0_fix_fp16.safetensors",
        variant="fp16",
        torch_dtype=torch.float16
    ).to("cuda")

    # LoRAウェイトの準備
    pipe.load_lora_weights("latent-consistency/lcm-lora-sdv1-5")

    # スケジューラの準備
    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
    # 画像を生成

    from diffusers.utils import load_image

    # プロンプトと画像の準備
    prompt = "プロンプト"
    image_url = "参照する画像のURL"

    import random

    random_seed = random.randint(-999, 999)

    # 画像生成
    image = pipe(
        prompt=prompt,
        image=load_image(image_url),
        num_inference_steps=4,
        guidance_scale=0,  # guide_scaleの無効化
        generator=torch.Generator("cpu").manual_seed(random_seed),
    ).images[0]

    import io
    import base64

    # バイトIOオブジェクトを初期化
    buf = io.BytesIO()
    # 画像をPNG形式でバイトIOオブジェクトに保存
    image.save(buf, format="PNG")
    # バイトIOオブジェクトをBase64エンコードする
    return base64.b64encode(buf.getvalue()).decode('utf-8')