在 colab 透過 Stable Diffusion 訓練, colab 需要下那些指令?

數據集存放在 Google Drive 的 /AI/datasets/zenmaru_dataset，輸出模型存放在 /AI/output/zenmaru_lora/，使用 Diffusers，不依賴 WebUI 或 Kohya-ss。

---

1. 檢查 GPU

先確認 Colab 是否有 GPU：

!nvidia-smi

如果沒有 GPU，請到 “執行階段” > “變更執行階段類型”，選擇 GPU。

---

2. 安裝 Diffusers 及其他依賴

!pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
!pip install diffusers transformers accelerate safetensors datasets peft

---

3. 掛載 Google Drive 並設置數據集與輸出路徑

將 Google Drive 掛載，並設置 數據集與 輸出模型存放目錄：

from google.colab import drive
drive.mount('/content/drive')

# 設定 Google Drive 路徑
DATASET_DIR = "/content/drive/MyDrive/AI/datasets/zenmaru_dataset"
OUTPUT_DIR = "/content/drive/MyDrive/AI/output/zenmaru_lora"

# 確保輸出目錄存在
import os
os.makedirs(OUTPUT_DIR, exist_ok=True)

# 確認數據集是否成功掛載
if os.path.exists(DATASET_DIR):
    print("✅ 數據集已成功掛載！")
    print("📂 數據集內容:", os.listdir(DATASET_DIR))
else:
    print("❌ 無法找到數據集，請檢查路徑是否正確！")

---

4. 下載 Stable Diffusion 2.1 預訓練模型

使用 Stable Diffusion 2.1：

from diffusers import StableDiffusionPipeline
import torch

PRETRAINED_MODEL_NAME = "stabilityai/stable-diffusion-2-1"

pipe = StableDiffusionPipeline.from_pretrained(PRETRAINED_MODEL_NAME, torch_dtype=torch.float16)
pipe.to("cuda")

如果你有自己的 .safetensors 模型：

!mkdir -p /content/models
%cd /content/models
!wget -c https://huggingface.co/stabilityai/stable-diffusion-2-1/raw/main/sd2-1_768_v-ema.safetensors -O model.safetensors

---

5. 訓練 LoRA

LoRA 訓練是微調模型的輕量方式：

from diffusers import UNet2DConditionModel
from peft import LoraConfig, get_peft_model
from transformers import Trainer, TrainingArguments
from datasets import load_dataset
from torchvision import transforms
from PIL import Image
import torch

# 設定 LoRA 參數
lora_config = LoraConfig(
    r=8,  # LoRA Rank（較低的值會節省顯存）
    lora_alpha=16,
    lora_dropout=0.1,
    target_modules=["attn1.to_q", "attn1.to_k", "attn1.to_v", # Replace with actual module names
                     "attn2.to_q", "attn2.to_k", "attn2.to_v"]  # 影響 Transformer 權重
)

# 載入 Stable Diffusion 的 UNet 模型
unet = UNet2DConditionModel.from_pretrained(
    PRETRAINED_MODEL_NAME, 
    subfolder="unet", 
    torch_dtype=torch.float16
)

# 將 UNet 模型包裝為 LoRA 模型
unet = get_peft_model(unet, lora_config)
unet.print_trainable_parameters()  # 顯示可訓練參數

# 訓練用的轉換函數（灰階轉 RGB，並新增一個通道）
def transform_gray_to_rgb(example):
    # Access image data using 'image' key, no need for as_posix()
    img = example["image"]  # example["image"] is already a PIL Image object
    img = img.convert("RGB")  # 轉為 RGB
    img = transforms.ToTensor()(img)  # 轉為 Tensor

    # Add a fourth channel filled with 1s (or other desired values)
    # Assume img shape is (3, H, W)
    # We want to add a channel to make it (4, H, W)
    fourth_channel = torch.ones(img.shape[1], img.shape[2], dtype=img.dtype, device=img.device) # Create a fourth channel filled with 1s
    img = torch.cat([img, fourth_channel.unsqueeze(0)], dim=0)  # Add the fourth channel to the image tensor

    example["pixel_values"] = img
    return example

# 載入數據集並進行灰階圖片動態轉換
dataset = load_dataset("imagefolder", data_dir=DATASET_DIR, split="train")
dataset = dataset.map(transform_gray_to_rgb)

# 設定訓練參數
training_args = TrainingArguments(
    output_dir=OUTPUT_DIR,
    per_device_train_batch_size=2,  # 根據 Colab GPU 記憶體調整
    gradient_accumulation_steps=4,  # 讓小 batch 也能訓練更久
    learning_rate=1e-4,  # LoRA 建議學習率
    max_steps=50,  # 訓練 5000 步
    save_strategy="steps",
    save_steps=10,  # 每 1000 步存一次
    report_to="none",
    #  Add this line to disable removing unused columns:
    remove_unused_columns=False  
)

# Custom data collator function
def custom_data_collator(features):
    pixel_values = [feature["pixel_values"] for feature in features]
    # Ensure all elements are tensors before stacking
    pixel_values = [torch.tensor(pv) if not isinstance(pv, torch.Tensor) else pv for pv in pixel_values]  # Convert to tensor if not already
    # Stack pixel values into a single tensor
    pixel_values = torch.stack(pixel_values)
    # Return a dictionary with the stacked pixel values, a dummy timestep and dummy encoder_hidden_states
    batch_size = pixel_values.shape[0]
    # Create random encoder hidden states with shape (batch_size, sequence_length, hidden_size)
    # Assuming sequence_length = 77 and hidden_size = 768 for Stable Diffusion 2.1
    encoder_hidden_states = torch.randn(batch_size, 77, 768, device=pixel_values.device, dtype=torch.float16) # Changed to float16 
    return {"sample": pixel_values.half(), "timestep": torch.tensor([1] * pixel_values.shape[0], device=pixel_values.device, dtype=torch.float16), "encoder_hidden_states": encoder_hidden_states}  # Changed sample and timestep to half and float16 respectively

# ... (training_args definition)

# 設定 Trainer
trainer = Trainer(
    model=unet,
    args=training_args,
    train_dataset=dataset,
    # Use the custom data collator
    data_collator=custom_data_collator 
)

# 開始 LoRA 訓練
trainer.train()

# 儲存 LoRA 模型
trainer.save_model(OUTPUT_DIR)
print(f"✅ LoRA 模型已儲存到 {OUTPUT_DIR}")

參數說明：

• DATASET_DIR = "/content/drive/MyDrive/AI/datasets/zenmaru_dataset"：數據集來自 Google Drive
• OUTPUT_DIR = "/content/drive/MyDrive/AI/output/zenmaru_lora"：訓練後的模型存放在 Google Drive
• max_steps=2000：訓練 2000 步
• per_device_train_batch_size=2：根據 GPU 記憶體調整

這段程式碼會進行以下步驟：

1. 設定 LoRA 參數並將 Stable Diffusion UNet 模型轉換為 LoRA 模型。
2. 動態將灰階圖片轉換為 RGB，並將其轉換為 Tensor。
3. 設定訓練參數，並使用 Trainer 開始訓練。
4. 訓練完成後，儲存 LoRA 模型到指定的輸出目錄。

---

6. 訓練完整模型（Fine-tuning）

如果你想完整微調 Stable Diffusion 2.1：

from diffusers import DDPMScheduler, UNet2DConditionModel
from transformers import TrainingArguments, Trainer
from datasets import load_dataset

# 載入 UNet 模型（Stable Diffusion 的關鍵部分）
unet = UNet2DConditionModel.from_pretrained(
    PRETRAINED_MODEL_NAME,
    subfolder="unet",
    torch_dtype=torch.float16
).to("cuda")

# 設定訓練 scheduler
pipe.scheduler = DDPMScheduler.from_config(pipe.scheduler.config)

# 設定訓練參數
training_args = TrainingArguments(
    output_dir=OUTPUT_DIR,  # 設定輸出目錄
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,  # 增加累積梯度，降低 GPU 記憶體需求
    learning_rate=5e-6,  # SD 2.1 fine-tuning 建議 5e-6
    max_steps=10000,  # 訓練 10,000 步
    lr_scheduler_type="constant",
    save_strategy="steps",
    save_steps=1000,  # 每 1000 步保存一次
    report_to="none"  # 避免 logging 到 wandb
)

# 載入數據集（需符合 Diffusers 格式）
dataset = load_dataset("imagefolder", data_dir=DATASET_DIR, split="train")

# 設定 Trainer
trainer = Trainer(
    model=unet,
    args=training_args,
    train_dataset=dataset
)


trainer.train()

Fine-tuning 需要更強的 GPU，建議使用 LoRA 替代。

---

7. 確認輸出結果並下載

訓練結束後，確認輸出結果：

import os
print("📂 訓練輸出內容:", os.listdir(OUTPUT_DIR))

如果要下載 LoRA 模型到本地：

from google.colab import files
files.download(f"{OUTPUT_DIR}/model.safetensors")

---

這樣你就能 使用 Diffusers 訓練 Stable Diffusion 2.1，
數據集來自 Google Drive 的 /AI/datasets/zenmaru_dataset，
輸出結果存放在 Google Drive 的 /AI/output/zenmaru_lora/，
不需要 WebUI 或 Kohya-ss！🚀

你打算訓練 LoRA 還是完整微調（Fine-tuning）呢？😃

---

注意事項：

• Google Colab 的免費版本有使用限制，長時間的訓練可能會被中斷。
• 使用付費的 Google Colab Pro 或 Pro+ 可以獲得更長的運行時間和更強的 GPU 資源。
• Stable Diffusion 的訓練需要大量的 GPU 記憶體，請確保你的 Colab 運行時有足夠的記憶體。
• LoRA訓練可以大幅降低對於硬體資源的需求，並且可以更快速的訓練。