見出し画像

ローカル環境のFlux + DiffusersでAIグラビア

ComfyUIからDiffusersへの移行

現在、ComfyUIからDiffusersへの乗り換えを進めています。
自然言語によるプログラミングができる「Cursor」がある今、GUIの方が簡単と言えなくなっていると思います。
今回は、FluxをDiffusers上で動作させることに挑戦しましたので、その経緯を共有します。

DiffusersでのFlux動作とモデル量子化

初めにオリジナルのFluxモデルを試みたものの、私の環境では動作しませんでした。

そのため、fp8形式に量子化を行い、モデルの軽量化を図りました。
プログラムに組み込みましたので、同じ課題に直面している方は参考にしていただければと思います。

image to image機能の実装について

次に、image to imageの機能を追加しようと試みました。しかし、現状ではFlux Pipelineがこの機能に対応していません。Flux Pipelineに依存しない独自の実装を模索し、GPT-4oやClaude-sonnetも併用して試みましたが、望んだ結果を得ることはできませんでした。

一方で、ComfyUIではFluxを用いたimage to image機能がすでにサポートされています。最新機能を試す場合は、ComfyUIを検討する価値があります。
しかし、ComfyUIはエラーが頻発し、環境を再構築するたびにワークシートが消失することが多く、私は今後もDiffusersを主軸に開発を進める予定です。

AIグラビア

先日、以下のオンラインセミナーに参加しました。
「AIによるグラビア画像生成」セミナーのレポートはこちら

セミナーでは、生成モデルを用いた創作手法が紹介されていました。
私も一度はグラビア画像を生成してみたいと思っていたので、セミナーで教わったプロンプトを利用して、実際に挑戦してみました。
しかし、生成される画像はどうしてもAI特有の質感が残ってしまい、自然な表現には一歩及びませんね。

AIグラビア

Koda Diffusionでの改善とfp8量子化の課題

そこで、Loraモデル「Koda Diffusion」を使用してみました。
Koda Diffusion (Flux)のダウンロードはこちら

このモデルを導入すると、AI特有の表現を抑えたより自然な画像が生成できるようになるようです。fp16形式でLoraを適用し、その後fp8形式に量子化を行いましたが、残念ながら最終的に真っ黒な画像が生成されてしまいました。fp8形式専用のLoraモデルでないと正しく動作しないようです。

私のPCのスペック

  • 環境: Windows 11 Home、WSL、Docker

  • プロセッサ: 12th Gen Intel(R) Core(TM) i5-12400 (2.50 GHz)

  • メモリ: 64.0 GB(使用可能: 63.8 GB)

  • GPU: NVIDIA GeForce RTX 4070 Ti SUPER 16GB

  • SSD: 1TB

実行環境とコードを載せておきますね!

Dockerfile

FROM nvcr.io/nvidia/pytorch:24.07-py3

RUN apt-get update && apt-get install -y wget
RUN apt-get install -y pulseaudio alsa-utils

RUN pip install --no-cache-dir openai langchain playwright lxml asyncio nest_asyncio beautifulsoup4

RUN playwright install-deps
RUN playwright install

ENV TZ='Asia/Tokyo'
WORKDIR /workspaces

devcontainer.json

{
  "name": "pytorch",
  "dockerFile": "Dockerfile",
  "runArgs": [
    "--gpus",
    "all",
    "--shm-size",
    "64gb",
    "-e",
    "PULSE_SERVER=unix:/mnt/wslg/PulseServer",
    "-v",
    "/mnt/wslg/:/mnt/wslg/"
  ],
  "settings": {
    "terminal.integrated.shell.linux": "/bin/bash"
  },
  "extensions": [
    "dbaeumer.vscode-eslint",
    "esbenp.prettier-vscode"
  ],
  "workspaceMount": "source=${localWorkspaceFolder},target=/hostdir,type=bind,consistency=cached",
  "workspaceFolder": "/workspaces",
  "containerEnv": {
    "MOUNTED_HOST_DIR": "${localWorkspaceFolder}",
    "MOUNTED_HOST_DIR_PATH_IN_CONTAINER": "/hostdir"
  },
  "forwardPorts": [80],
  "features": {
    "docker-in-docker": {
      "version": "latest"
    }
  },
  "remoteUser": "root",
  "containerUser": "root"
}

コンテナ作成後の実行コマンド

pip install --upgrade pip 
pip install transformers sentencepiece accelerate diffusers optimum-quanto huggingface_hub
pip uninstall torchvision flash_attn
pip install torchvision flash_attn
pip uninstall apex
pip install apex

ImageModelHandler.py

import os
import torch
import json
from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler
from diffusers.models.transformers.transformer_flux import FluxTransformer2DModel
from diffusers.pipelines.flux.pipeline_flux import FluxPipeline
from optimum.quanto import freeze, qfloat8, quantize, quantization_map
from optimum.quanto import QuantizedTransformersModel, QuantizedDiffusersModel
from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
from huggingface_hub import login
from PIL import Image

class ImageModelHandler:
  def __init__(self, model_version="FLUX.1-dev"):
    """
    コンストラクタ: モデルのバージョン、データ型、ディレクトリパスを設定
    model_version: 使用するモデルのバージョンを指定
    dtype: モデル読み込み時に使用するデータ型を指定(torch.bfloat16 推奨)
    """
    self.model_version = model_version
    self.dtype = torch.bfloat16
    self.quantized_model_dir = f"./{model_version}-distilled"
    self.hf_token = os.getenv('HUGGINGFACE_TOKEN')
    self.bfl_repo = f"black-forest-labs/{model_version}"
    self._initialize_models()

  def _initialize_models(self):
    """ 各種サブモデル(スケジューラー、テキストエンコーダ、トークナイザ、VAE、量子化モデルなど)の初期化を行う """
    if not self.hf_token:
      raise ValueError("Hugging Faceのトークンが設定されていません。")
    login(token=self.hf_token)

    # スケジューラーの初期化: 画像生成プロセスを制御
    self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
      pretrained_model_name_or_path=self.bfl_repo,
      subfolder="scheduler",
    )

    # CLIP テキストエンコーダ: テキストプロンプトをベクトル形式に変換する(画像生成における初期条件)
    self.text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=self.dtype)

    # CLIP トークナイザ: テキストをトークンに変換し、テキストエンコーダへ渡す
    self.tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=self.dtype)

    # T5 トークナイザ: 高度なテキスト処理(サブプロンプトなど)に使用される補助トークナイザ
    self.tokenizer_2 = T5TokenizerFast.from_pretrained(self.bfl_repo, subfolder="tokenizer_2")

    # VAE(変分オートエンコーダ): 画像の圧縮と復元を行う(画像のクオリティを決定)
    self.vae = AutoencoderKL.from_pretrained(self.bfl_repo, subfolder="vae", torch_dtype=self.dtype)

    # Transformer と T5 エンコーダのパス設定
    transformer_path, text_encoder_2_path = self._get_model_paths()
    
    # Transformer および Text Encoder 2の読み込みまたは量子化
    self.transformer, self.text_encoder_2 = self._load_or_quantize_models(transformer_path, text_encoder_2_path)

    # 各サブモデルを統合し、FluxPipelineを構築
    self.pipe = FluxPipeline(
      scheduler=self.scheduler,
      text_encoder=self.text_encoder,
      tokenizer=self.tokenizer,
      text_encoder_2=self.text_encoder_2,
      tokenizer_2=self.tokenizer_2,
      vae=self.vae,
      transformer=self.transformer,
    )

    # メモリ節約のためのモデルオフロード設定
    self.pipe.enable_model_cpu_offload()

  def _get_model_paths(self):
    """ 量子化されたモデルの保存ディレクトリを指定 """
    return (
      os.path.join(self.quantized_model_dir, "transformer"),
      os.path.join(self.quantized_model_dir, "text_encoder_2")
    )

  def _load_or_quantize_models(self, transformer_path, text_encoder_2_path):
    """
    モデルを読み込み、量子化されていない場合は量子化を行う
    transformer_path, text_encoder_2_path: モデルの保存パスを指定
    """
    if os.path.exists(transformer_path) and os.path.exists(text_encoder_2_path):
      print("量子化されたモデルが見つかりました。量子化モデルを読み込みます。")

      # 量子化されたモデルを読み込む際にカスタムクラスを使用
      class QuantizedT5EncoderModelForCausalLM(QuantizedTransformersModel):
        auto_class = T5EncoderModel

      class QuantizedFluxTransformer2DModel(QuantizedDiffusersModel):
        base_class = FluxTransformer2DModel

      transformer = QuantizedFluxTransformer2DModel.from_pretrained(transformer_path).to(dtype=self.dtype)
      T5EncoderModel.from_config = lambda c: T5EncoderModel(c)
      text_encoder_2 = QuantizedT5EncoderModelForCausalLM.from_pretrained(text_encoder_2_path).to(dtype=self.dtype)
    else:
      print("量子化されたモデルが見つかりません。フル精度のモデルを読み込み、量子化します。")

      # 量子化されていない場合、フル精度のモデルを読み込み、量子化する
      text_encoder_2 = T5EncoderModel.from_pretrained(self.bfl_repo, subfolder="text_encoder_2", torch_dtype=self.dtype)
      transformer = FluxTransformer2DModel.from_pretrained(self.bfl_repo, subfolder="transformer", torch_dtype=self.dtype)
      self._quantize_and_save(transformer, transformer_path)
      self._quantize_and_save(text_encoder_2, text_encoder_2_path)
    return transformer, text_encoder_2

  def _quantize_and_save(self, model, save_path):
    """
    モデルを量子化し、指定されたパスに保存する
    model: 量子化するモデルオブジェクト
    save_path: 保存先のディレクトリ
    """
    print(f"モデルを量子化して保存します: {save_path}")
    quantize(model, weights=qfloat8)
    freeze(model)
    os.makedirs(save_path, exist_ok=True)
    model.save_pretrained(save_path)
    with open(os.path.join(save_path, "quanto_qmap.json"), "w") as f:
      json.dump(quantization_map(model), f)
    print(f"量子化されたモデルを保存しました: {save_path}")

  def text_to_image(self, prompt, width=768, height=768, num_inference_steps=25, guidance_scale=7.5, seed=None, max_sequence_length=256, output_path="output_image.png"):
    """
    テキストプロンプトを用いて画像を生成する(パラメータは一般的に推奨される値に設定)
    prompt: 生成したい画像のテキストプロンプト
    width, height: 生成画像のサイズ(幅と高さをピクセル単位で指定)
    num_inference_steps: 生成過程におけるステップ数(多いほど高品質だが時間がかかる)
    guidance_scale: テキストプロンプトに対する忠実度(値が高いほどプロンプトに従うが、自然さが失われる場合も)
    seed: 生成のランダムシード(再現性のために指定)
    max_sequence_length: プロンプトの最大長(長すぎるとエラーになることがある)
    output_path: 生成された画像の保存パス
    """
    if seed is not None:
      generator = torch.Generator().manual_seed(seed)
    else:
      generator = torch.Generator()

    print(f"画像生成を開始します。プロンプト: '{prompt}'")
    image = self.pipe(
      prompt=prompt,
      width=width,
      height=height,
      num_inference_steps=num_inference_steps,
      max_sequence_length=max_sequence_length,
      generator=generator,
      guidance_scale=guidance_scale,
    ).images[0]

    image.save(output_path)
    print(f"生成された画像を保存しました: {output_path}")
    return image

# メイン関数の実行(モデルの選択と画像生成)
if __name__ == "__main__":
  available_models = {
    "default": "FLUX.1-dev",
    "fast": "FLUX.1-schnell"
  }
  selected_model_version = available_models["default"]
  print(f"選択されたモデルバージョン: {selected_model_version}")

  handler = ImageModelHandler(model_version=selected_model_version)

  # テスト画像の生成
  handler.text_to_image(
    prompt="A futuristic cityscape at sunset with glowing skyscrapers",  # 描写したい画像の内容をテキストで指定
    width=512,  # 幅を指定(一般的には512または768が推奨される)
    height=512,  # 高さを指定(同じく512または768)
    num_inference_steps=30,  # ステップ数を指定(ステップ数が多いほど高品質になるが、処理時間も増加)
    guidance_scale=7.5,  # ガイダンススケールを指定(7.5は一般的な設定、低すぎるとプロンプトに従わない画像が生成される)
    seed=None,  # ランダムシードを固定(再現性を確保したい場合に使用)None を設定することで毎回異なる画像を生成する
    max_sequence_length=256,  # テキストプロンプトの最大長(256は通常の設定)
    output_path="generated_image.png"  # 生成された画像の保存パス
  )


プログラミングに苦手意識をお持ちの方は、ぜひ「DXパーティー」までお問い合わせください。
有料サービスではありますが、オンラインでマンツーマンのレッスンを受けることで、短期間でPythonを習得し、自在に使いこなせるようになると思います!

お問い合わせ - DXパティー (dxpt.jp)

マンツーマンオンライン講師を募集しています。 特に視聴者数が少ないDX・デジタル領域に取り組んでいるYouTube配信者を積極的に募集しています。 DXパティーのレッスンとYouTubeチャンネルの収益化の両方から収益を得て、安定した収入を目指しましょう。

DXパティーでは講師募集中!

お仕事のご依頼もお待ちしております!