これは「Hunyuan Video」ベースのLoRAをColabローカル上で作成するためのノートブックとなります。学習素材は画像・動画のいずれかが利用可能です。静止画像5枚から問題なくHunyuanVideo LoRAが作れることを確認しています。使い方は各コードセル部分の説明をご覧ください。

hunyuanvideo_lora_training.sh

### ---この区切りごとにコピーして、それぞれ別のコードセルに貼り付けていってください---

#@title # HunyuanVideo LoRA作成ノートブック

#@markdown ---

#@markdown ## このノートブックについて

#@markdown これは「Hunyuan Video」ベースのLoRAをColabローカル上で作成するためのノートブックとなります。

#@markdown [jhj0517/finetuning-notebooks](https://github.com/jhj0517/finetuning-notebooks)をベースに作成しています。

#@markdown Google Driveにはマウントしない作りとなっているため、学習途中や完了後に得られるファイルは**ランタイム再起動時に消去**されます。ご注意ください。

#@markdown 必要な成果物(LoRA, checkpointなど)を保存したい場合は、**都度手動ダウンロード**するか、各セル実行後に別途Google Drive等にアップロードするなどして管理を行ってください。

#@markdown もしGoogle Driveにマウントして使いたい場合、ベースとなっている[jhj0517/finetuning-notebooks](https://github.com/jhj0517/finetuning-notebooks)をご利用ください。

### ---この区切りごとにコピーして、それぞれ別のコードセルに貼り付けていってください---

#@title ## 1. diffusion-pipe、ならびに必要なライブラリのインストール
#@markdown This notebook is powered by https://github.com/tdrussell/diffusion-pipe

!git clone --recurse-submodules https://github.com/tdrussell/diffusion-pipe
%cd diffusion-pipe

!pip install deepspeed
!pip install datasets
!pip install torch-optimi
!pip install bitsandbytes
!pip install av
!pip install loguru
!pip install flash-attn

### ---この区切りごとにコピーして、それぞれ別のコードセルに貼り付けていってください---

#@title ## ２．データセットの追加

#@markdown このコードセル内のコマンドを実行すると`/content/dataset` というフォルダが作成されます。

#@markdown （Colab 左側の「ファイル」タブから確認できます）

#@markdown そこへ**画像や動画、およびそれぞれ対応する `.txt` ファイル**をアップロードします。

#@markdown 例:

#@markdown /content/dataset/ ├── file1.png ├── file1.txt # "file1.png"に対応するプロンプト ├── file2.mp4 ├── file2.txt # "file2.mp4"に対応するプロンプト ├── ...

#@markdown 各動画ファイルや画像ファイルと同名の`.txt`を用意し、テキスト内容に「画像や動画のプロンプト」を書いておいてください。

#@markdown もし、フレーム数の異なる動画が混在する場合には `frame_buckets` でそのフレーム数をすべて列挙しておくことを推奨します

!mkdir /content/dataset

### ---この区切りごとにコピーして、それぞれ別のコードセルに貼り付けていってください---

#@title ## ３．キャプションファイルの作成（すでにキャプションファイルがある人は不要です）

#@markdown `.txt` が存在しない画像や動画ファイルについて、デフォルトのプロンプトで `.txt` ファイルを自動生成します。

#@markdown 画像と動画で別々のプロンプトを指定することもできます。

#@markdown（ `default_xxxxx_prompt` にプロンプトをいれることで、`dataset` フォルダ内に格納したすべての画像に対してキャプションファイルを作成します。なお、ファイルによってキャプションを変えたい場合は手作業で修正してください。）

import os
from glob import glob

DATASET_PATH = "/content/dataset"  # @param {type:"string"}

default_image_prompt = "a image of something"  # @param {type:"string"}
default_video_prompt = "a video of something"  # @param {type:"string"}

IMAGE_EXTS = [".png", ".jpg", ".jpeg", ".webp", ".bmp"]
VIDEO_EXTS = [".mp4", ".webm", ".m4v", ".mov", ".avi"]

os.makedirs(DATASET_PATH, exist_ok=True)

def create_txt_if_not_exists(file_path, prompt_text):
    base, _ = os.path.splitext(file_path)
    txt_path = base + ".txt"
    if not os.path.exists(txt_path):
        with open(txt_path, "w") as f:
            f.write(prompt_text)
        print(f"Created: {txt_path}")
    else:
        # .txt が既に存在する場合は何もしない
        pass

all_files = glob(os.path.join(DATASET_PATH, "*"))

for f in all_files:
    ext = os.path.splitext(f)[1].lower()
    if ext in IMAGE_EXTS:
        create_txt_if_not_exists(f, default_image_prompt)
    elif ext in VIDEO_EXTS:
        create_txt_if_not_exists(f, default_video_prompt)

print("Done generating missing .txt files.")

### ---この区切りごとにコピーして、それぞれ別のコードセルに貼り付けていってください---

#@title # 4. モデルのダウンロード (Hunyuan, CLIP, LLM)

#@markdown このセルを実行すると、選択したモデルがColabローカル(`"/content/hunyuan/base_models"`)にダウンロードされます。

#@markdown （変更の必要がなければ、デフォルト設定のままでも問題ありません）

import huggingface_hub
import os

# Set local base directory for downloaded models
BASE_MODELS_DIR_PATH = "/content/hunyuan/base_models" # @param {type:"string"}

#@markdown ダウンロードするベースモデルを選択します
HUNYUAN_VIDEO = "hunyuan_video_720_cfgdistill_fp8_e4m3fn.safetensors" #@param ["hunyuan_video_720_cfgdistill_fp8_e4m3fn.safetensors", "hunyuan_video_720_cfgdistill_bf16.safetensors", "hunyuan_video_FastVideo_720_fp8_e4m3fn.safetensors"]
VAE           = "hunyuan_video_vae_bf16.safetensors"                 #@param ["hunyuan_video_vae_bf16.safetensors", "hunyuan_video_vae_fp32.safetensors"]
CLIP          = "openai/clip-vit-large-patch14"                      #@param ["openai/clip-vit-large-patch14"]
LLM           = "Kijai/llava-llama-3-8b-text-encoder-tokenizer"      #@param ["Kijai/llava-llama-3-8b-text-encoder-tokenizer"]

os.makedirs(BASE_MODELS_DIR_PATH, exist_ok=True)
HUNYUAN_MODELS_DIR = os.path.join(BASE_MODELS_DIR_PATH, "hunyuan")
os.makedirs(HUNYUAN_MODELS_DIR, exist_ok=True)

CLIP_MODEL_DIR = os.path.join(BASE_MODELS_DIR_PATH, "clip")
os.makedirs(CLIP_MODEL_DIR, exist_ok=True)

LLM_MODEL_DIR = os.path.join(BASE_MODELS_DIR_PATH, "llm")
os.makedirs(LLM_MODEL_DIR, exist_ok=True)

HUNYUAN_VIDEO_URL = {
    "hunyuan_video_720_cfgdistill_fp8_e4m3fn.safetensors": "https://huggingface.co/Kijai/HunyuanVideo_comfy/resolve/main/hunyuan_video_720_cfgdistill_fp8_e4m3fn.safetensors",
    "hunyuan_video_720_cfgdistill_bf16.safetensors":       "https://huggingface.co/Kijai/HunyuanVideo_comfy/resolve/main/hunyuan_video_720_cfgdistill_bf16.safetensors",
    "hunyuan_video_FastVideo_720_fp8_e4m3fn.safetensors":  "https://huggingface.co/Kijai/HunyuanVideo_comfy/resolve/main/hunyuan_video_FastVideo_720_fp8_e4m3fn.safetensors",
}

HUNYUAN_VIDEO_VAE_URL = {
    "hunyuan_video_vae_bf16.safetensors": "https://huggingface.co/Kijai/HunyuanVideo_comfy/resolve/main/hunyuan_video_vae_bf16.safetensors",
    "hunyuan_video_vae_fp32.safetensors": "https://huggingface.co/Kijai/HunyuanVideo_comfy/resolve/main/hunyuan_video_vae_fp32.safetensors",
}

video_url = HUNYUAN_VIDEO_URL[HUNYUAN_VIDEO]
vae_url   = HUNYUAN_VIDEO_VAE_URL[VAE]
clip_repo_id = CLIP
llm_repo_id  = LLM

print("Downloading Hunyuan base model to:", HUNYUAN_MODELS_DIR)
!wget {video_url} -P {HUNYUAN_MODELS_DIR}
print("Downloading VAE to:", HUNYUAN_MODELS_DIR)
!wget {vae_url}   -P {HUNYUAN_MODELS_DIR}

print("Downloading CLIP to:", CLIP_MODEL_DIR)
huggingface_hub.snapshot_download(clip_repo_id, local_dir=CLIP_MODEL_DIR)

print("Downloading LLM to:", LLM_MODEL_DIR)
huggingface_hub.snapshot_download(llm_repo_id, local_dir=LLM_MODEL_DIR)

print("Done.")

### ---この区切りごとにコピーして、それぞれ別のコードセルに貼り付けていってください---

#@title # 5. トレーニング処理とパラメーター

#@markdown データセットとして画像を与えて学習させる場合、このままの設定で学習が可能です

#@markdown 学習が完了したLoRAファイルは、 `/content/hunyuan/outputs` 配下に出力されます

import os
import sys
import toml

#@markdown LoRA学習の際、以前のチェックポイントからトレーニングを再開させたい場合は、こちらをチェックしてください。
RESUME_FROM_CHECKPOINT = False #@param {type:"boolean"}

#@markdown ## 各ファイルのパス設定
BASE_MODELS_DIR_PATH = "/content/hunyuan/base_models"   # @param {type:"string"}
OUTPUT_LORA_NAME     = "My-Hunyuan-Lora-V1"             # @param {type:"string"}
OUTPUT_LORA_DIR_PATH = "/content/hunyuan/outputs"       # @param {type:"string"}
DATASET_PATH         = "/content/dataset"               # @param {type:"string"}

OUTPUT_LORA_DIR_PATH = os.path.join(OUTPUT_LORA_DIR_PATH, OUTPUT_LORA_NAME)
transformer_dir_path = os.path.join(BASE_MODELS_DIR_PATH, "hunyuan")
vae_dir_path         = os.path.join(BASE_MODELS_DIR_PATH, "hunyuan")
clip_path            = os.path.join(BASE_MODELS_DIR_PATH, "clip")
llm_path             = os.path.join(BASE_MODELS_DIR_PATH, "llm")

os.makedirs(OUTPUT_LORA_DIR_PATH, exist_ok=True)
os.makedirs(DATASET_PATH,         exist_ok=True)

#@markdown ## データセットの構成（画像で学習させる際はこのまま実行できます）
#@markdown - `frame_buckets` は、データセットのフレーム番号のリストです。
#@markdown   - 例えば、データセットに30フレーム、60フレームの動画が含まれている場合は、 `[30, 60]` を使います。
#@markdown   - データセットに画像も含まれている場合は、 `[1]` も含める。
#@markdown - `resolutions` は、 `diffusion-pipe` がデータセットをリサイズする解像度のリストです
#@markdown   - VRAMが24GB以下の場合は、512以下に設定してください。

frame_buckets = [1]   # @param {type:"raw"}
resolutions   = [512] # @param {type:"raw"}

enable_ar_bucket = True  # @param {type:"boolean"}
min_ar = 0.5             # @param {type:"number"}
max_ar = 2.0             # @param {type:"number"}
num_ar_buckets = 7       # @param {type:"integer"}

# Repetition factor for each dataset entry
num_repeats = 5

# Write dataset.toml
dataset_config = {
    "resolutions": resolutions,
    "frame_buckets": frame_buckets,
    "enable_ar_bucket": enable_ar_bucket,
    "min_ar": min_ar,
    "max_ar": max_ar,
    "num_ar_buckets": num_ar_buckets,
    "directory": [
        {
            "path": DATASET_PATH,
            "num_repeats": num_repeats,
        }
    ],
}

dataset_config_file_path = os.path.join(DATASET_PATH, "dataset.toml")
with open(dataset_config_file_path, "w") as toml_file:
    toml.dump(dataset_config, toml_file)

print(f"dataset.toml is saved to {dataset_config_file_path}")

#@markdown ## Base Model Configuration
BASE_HUNYUAN_MODEL_NAME = "hunyuan_video_720_cfgdistill_fp8_e4m3fn.safetensors" #@param ["hunyuan_video_720_cfgdistill_fp8_e4m3fn.safetensors", "hunyuan_video_720_cfgdistill_bf16.safetensors", "hunyuan_video_FastVideo_720_fp8_e4m3fn.safetensors"]
BASE_VAE_MODEL_NAME      = "hunyuan_video_vae_bf16.safetensors"                 #@param ["hunyuan_video_vae_bf16.safetensors", "hunyuan_video_vae_fp32.safetensors"]

transformer_path = os.path.join(transformer_dir_path, BASE_HUNYUAN_MODEL_NAME)
vae_path         = os.path.join(vae_dir_path, BASE_VAE_MODEL_NAME)

model_type            = 'hunyuan-video'
dtype                 = 'bfloat16'         # @param {type:"string"}
transformer_dtype     = 'float8'           # @param {type:"string"}
timestep_sample_method= 'logit_normal'     # @param {type:"string"}

#@markdown ## Training Settings
epochs                     = 50  # @param {type:"integer"}
micro_batch_size_per_gpu   = 1   # @param {type:"integer"}
pipeline_stages            = 1   # @param {type:"integer"}
gradient_accumulation_steps= 4   # @param {type:"integer"}
gradient_clipping          = 1.0 # @param {type:"number"}
warmup_steps               = 50  # @param {type:"integer"}

#@markdown ## Eval Settings
eval_every_n_epochs               = 1    # @param {type:"integer"}
eval_before_first_step            = True # @param {type:"boolean"}
eval_micro_batch_size_per_gpu     = 1    # @param {type:"integer"}
eval_gradient_accumulation_steps  = 1    # @param {type:"integer"}

#@markdown ## LoRA Settings
adapter_type  = 'lora'
rank          = 32              # @param {type:"integer"}
adapter_dtype = 'bfloat16'      # @param {type:"string"}

#@markdown ## Optimizer Settings
optimizer_type = 'adamw_optimi' # @param {type:"string"}
lr             = 2e-5          # @param {type:"number"}
betas          = [0.9, 0.99]   # @param {type:"raw"}
weight_decay   = 0.02          # @param {type:"number"}
eps            = 1e-8          # @param {type:"number"}

#@markdown ## Misc Settings
save_every_n_epochs       = 10            # @param {type:"integer"}
checkpoint_every_n_minutes= 30            # @param {type:"integer"}
activation_checkpointing  = True          # @param {type:"boolean"}
partition_method          = 'parameters'  # @param {type:"string"}
save_dtype                = 'bfloat16'    # @param {type:"string"}
caching_batch_size        = 1             # @param {type:"integer"}
steps_per_print           = 1             # @param {type:"integer"}
video_clip_mode           = 'single_middle' # @param {type:"string"}

# Write config.toml
train_config = {
    "output_dir": OUTPUT_LORA_DIR_PATH,
    "dataset": dataset_config_file_path,

    "epochs": epochs,
    "micro_batch_size_per_gpu": micro_batch_size_per_gpu,
    "pipeline_stages": pipeline_stages,
    "gradient_accumulation_steps": gradient_accumulation_steps,
    "gradient_clipping": gradient_clipping,
    "warmup_steps": warmup_steps,

    "eval_every_n_epochs": eval_every_n_epochs,
    "eval_before_first_step": eval_before_first_step,
    "eval_micro_batch_size_per_gpu": eval_micro_batch_size_per_gpu,
    "eval_gradient_accumulation_steps": eval_gradient_accumulation_steps,

    "save_every_n_epochs": save_every_n_epochs,
    "checkpoint_every_n_minutes": checkpoint_every_n_minutes,
    "activation_checkpointing": activation_checkpointing,
    "partition_method": partition_method,
    "save_dtype": save_dtype,
    "caching_batch_size": caching_batch_size,
    "steps_per_print": steps_per_print,
    "video_clip_mode": video_clip_mode,

    "model": {
        "type": model_type,
        "transformer_path": transformer_path,
        "vae_path": vae_path,
        "llm_path": llm_path,
        "clip_path": clip_path,
        "dtype": dtype,
        "transformer_dtype": transformer_dtype,
        "timestep_sample_method": timestep_sample_method,
    },

    "adapter": {
        "type": adapter_type,
        "rank": rank,
        "dtype": adapter_dtype,
    },

    "optimizer": {
        "type": optimizer_type,
        "lr": lr,
        "betas": betas,
        "weight_decay": weight_decay,
        "eps": eps,
    },
}

train_config_file_path = os.path.join(DATASET_PATH, "config.toml")
with open(train_config_file_path, "w") as toml_file:
    toml.dump(train_config, toml_file)

print(f"config.toml is saved to {train_config_file_path}")

# Prepare environment for training
os.environ["NCCL_P2P_DISABLE"] = "1"
os.environ["NCCL_IB_DISABLE"]  = "1"

# Run training
if RESUME_FROM_CHECKPOINT:
  !deepspeed --num_gpus=1 train.py --deepspeed --config {train_config_file_path} --resume_from_checkpoint
else:
  !deepspeed --num_gpus=1 train.py --deepspeed --config {train_config_file_path}