Textual Inversion

app.py

@@ -7,10 +7,13 @@ from PIL import Image
 from streamlit_drawable_canvas import st_canvas
 
 from models.diffusion import Diffusion
+from train import finetune
 from utils import initialize_data_dir, save_image
 
 # 設定
-st.set_page_config(page_title="手書き文字生成アプリ", layout="wide")
+st.set_page_config(
+    page_title="手書き文字生成アプリ", layout="wide", page_icon=":pencil:"
+)
 
 # タイトル
 st.title("手書き文字生成アプリ")
@@ -32,7 +35,7 @@ num_samples = 5  # 各文字のサンプル数
 st.header("手書き文字を描いてください")
 
 # 保存用ディレクトリ
-data_dir = initialize_data_dir()
+data_dir = initialize_data_dir("./sample_images")
 
 # 描画領域の作成
 for char in characters:
@@ -53,33 +56,42 @@ for char in characters:
             if canvas.image_data is not None:
                 img = Image.fromarray(
                     canvas.image_data.astype("uint8"), "RGBA"
-                ).convert("L")
-                # 二値化
-                img = img.point(lambda x: 0 if x < 128 else 255, "1")
+                ).convert("RGB")
+                # 輝度を逆転
+                img = Image.eval(img, lambda x: 255 - x)
                 save_image(img, char, i, data_dir)
 
 # ハイパーパラメータの入力
-st.sidebar.header("ハイパーパラメータ設定")
+st.sidebar.header("学習パラメータ設定")
 learning_rate = st.sidebar.number_input(
     "学習率", min_value=0.0001, max_value=1.0, value=0.001, step=0.0001, format="%.4f"
 )
 epochs = st.sidebar.number_input(
     "エポック数", min_value=1, max_value=100, value=10, step=1
 )
-optimizer_name = st.sidebar.selectbox("最適化手法", ["SGD", "Adam", "RMSprop"])
-
+optimizer_name = st.sidebar.selectbox("最適化手法", ["SGD", "Adam", "RMSprop", "AdamW"])
+if optimizer_name == "SGD":
+    optimizer = torch.optim.SGD
+elif optimizer_name == "Adam":
+    optimizer = torch.optim.Adam
+elif optimizer_name == "RMSprop":
+    optimizer = torch.optim.RMSprop
+elif optimizer_name == "AdamW":
+    optimizer = torch.optim.AdamW
+num_augmentations = st.sidebar.number_input(
+    "データ拡張回数", min_value=0, max_value=100, value=20, step=1
+)
 
 # サンプリングの設定
 st.sidebar.header("サンプリング設定")
 noise_steps = st.sidebar.number_input(
-    "ノイズステップ数", min_value=1, max_value=1000, value=1000, step=1
-)
-beta_start = st.sidebar.number_input(
-    "βの初期値", min_value=0.0, max_value=1.0, value=0.0001, step=0.0001, format="%.4f"
+    "ノイズステップ数", min_value=2, max_value=1000, value=1000, step=1
 )
-beta_end = st.sidebar.number_input(
-    "βの終了値", min_value=0.0, max_value=1.0, value=0.02, step=0.0001, format="%.4f"
+num_chars = st.sidebar.number_input(
+    "生成文字数", min_value=1, max_value=46, value=5, step=1
 )
+beta_start = 0.0001
+beta_end = 0.02
 
 
 # 生成ボタン
@@ -107,26 +119,39 @@ if st.button("生成"):
             torch.load(save_path, weights_only=True, map_location=device)
         )
 
+        # ファインチューニング
+        progress_bar = st.progress(0, text="学習中...")
+        diffusion_model = finetune(
+            data_dir,
+            diffusion_model,
+            criterion=torch.nn.MSELoss(),
+            optimizer=torch.optim.AdamW,
+            num_epochs=epochs,
+            learning_rate=learning_rate,
+            num_augmentations=num_augmentations,
+            progress_bar=progress_bar,
+        )
+        progress_bar.empty()
+
         # フォント画像の生成と表示
         def chuncked(iterable, n):
             for i in range(0, len(iterable), n):
                 yield iterable[i : i + n]
 
-        labels = list(range(46))
+        labels = list(range(num_chars))
         columns_per_row = 5
         start_time = time.time()
         with st.spinner("フォント画像を生成中..."):
             labels_tensor = torch.tensor(labels).to(device)
-            font_image = diffusion_model.sample(
-                diffusion_model.model, labels_tensor
-            )
+            font_image = diffusion_model.sample(diffusion_model.model, labels_tensor)
             elapsed_time = time.time() - start_time
             st.success(f"フォント画像の生成に成功しました（{elapsed_time:.2f}秒）")
             for label_row in chuncked(labels, columns_per_row):
                 cols = st.columns(columns_per_row)
                 for col, label in zip(cols, label_row):
+                    # 輝度を逆転させて表示
                     col.image(
-                        font_image[label].permute(1, 2, 0).cpu().numpy(),
+                        255 - font_image[label].permute(1, 2, 0).cpu().numpy(),
                         caption=f"{label}",
                         use_container_width=True,
                     )

models/diffusion.py

@@ -1,5 +1,7 @@
 import numpy as np
 import torch
+from streamlit.delta_generator import DeltaGenerator
+
 from .unet import UNet_conditional
 
 
@@ -49,7 +51,9 @@ class Diffusion:
         ノイズ画像とノイズの生成
         """
         sqrt_alpha_hat = torch.sqrt(self.alpha_hat[t])[:, None, None, None]
-        sqrt_one_minus_alpha_hat = torch.sqrt(1 - self.alpha_hat[t])[:, None, None, None]
+        sqrt_one_minus_alpha_hat = torch.sqrt(1 - self.alpha_hat[t])[
+            :, None, None, None
+        ]
         noise = torch.randn_like(x)
         return sqrt_alpha_hat * x + sqrt_one_minus_alpha_hat * noise, noise
 
@@ -194,3 +198,64 @@ class Diffusion:
             history = np.vstack([history, item])
 
         return history
+
+    def fit_s(
+        self,
+        criterion: torch.nn.Module,
+        optimizer: torch.optim.Optimizer,
+        num_epochs: int,
+        learning_rate: float,
+        data_loader: torch.utils.data.DataLoader,
+        device: torch.device,
+        progress_bar: DeltaGenerator,
+    ) -> None:
+        """
+        Textual Inversionの学習
+        """
+        # 最小損失の初期化
+        min_train_loss = 9e9
+
+        # sのみを学習可能なパラメータとして定義
+        for params in self.model.parameters():
+            params.requires_grad = False
+        self.model.s = torch.nn.Parameter(
+            1e-2 * torch.randn(1, self.time_dim, device=device)
+        )
+        print("sの初期値: ", self.model.s)
+        optimizer = optimizer([self.model.s], lr=learning_rate)
+        for epoch in range(num_epochs):
+            print(f"Epoch {epoch + 1} / {num_epochs}")
+            train_loss = 0
+            n_train = 0
+
+            self.model.eval()
+            for x, labels in data_loader:
+                train_batch_size = len(labels)
+                n_train += train_batch_size
+
+                x = x.to(device)
+                labels = labels.to(device)
+
+                t = self.sample_timesteps(x.size(0)).to(device)
+                xt, noise = self.noise_images(x, t)
+
+                optimizer.zero_grad()
+                predicted_noise = self.model(xt, t, labels, self.model.s)
+                loss = criterion(predicted_noise, noise)
+                loss.backward()
+                optimizer.step()
+
+                train_loss += loss.item() * train_batch_size
+
+            # 損失計算
+            avg_train_loss = train_loss / n_train
+
+            # 最小損失の更新とsの保存
+            if avg_train_loss < min_train_loss:
+                min_train_loss = avg_train_loss
+
+            # 結果表示
+            print(f"Epoch {epoch + 1}, Train loss: {avg_train_loss:.3f}")
+            progress_bar.progress((epoch + 1) / num_epochs, text=f"学習中... (train loss = {avg_train_loss:.3f})")
+
+        print("sの最終値: ", self.model.s)

models/unet.py

@@ -174,14 +174,14 @@ class UNet_conditional(UNet):
         super().__init__(c_in, c_out, time_dim, **kwargs)
         self.label_emb = nn.Embedding(num_classes, time_dim)
 
-    def forward(self, x, t, y):
-        """
-        ラベルの埋め込みをタイムステップの埋め込みに加算
-        """
+    # ラベルの埋め込みとPsudo Wordをタイムステップの埋め込みに加算
+    def forward(self, x, t, labels=None, s=None):
         t = t.unsqueeze(-1)
         t = self.pos_encoding(t, self.time_dim)
 
-        if y is not None:
-            t += self.label_emb(y)
+        if labels is not None:
+            t += self.label_emb(labels)
+        if s is not None:
+            t += s
 
         return self.unet_forwad(x, t)

train.py

@@ -0,0 +1,61 @@
+from streamlit.delta_generator import DeltaGenerator
+import torch
+from torch.utils.data import DataLoader
+from torchvision import datasets, transforms
+
+from models.diffusion import Diffusion
+
+
+def finetune(
+        data_dir: str,
+        model: Diffusion,
+        criterion: torch.nn.Module,
+        optimizer: torch.optim.Optimizer,
+        num_epochs: int,
+        learning_rate: float,
+        num_augmentations: int,
+        progress_bar: DeltaGenerator,
+) -> Diffusion:
+    # transformの定義
+    transform = transforms.Compose(
+        [
+            transforms.Grayscale(num_output_channels=3),
+            transforms.Resize((32, 32)),
+            transforms.ToTensor(),
+            transforms.Normalize(0.5, 0.5),
+        ]
+    )
+
+    transform_aug = transforms.Compose(
+        [
+            transforms.Grayscale(num_output_channels=3),
+            transforms.Resize((32, 32)),
+            transforms.RandomAffine(
+                degrees=0, translate=(0.05, 0.05), scale=(0.95, 1.05), fill=255
+            ),
+            transforms.ToTensor(),
+            transforms.Normalize(0.5, 0.5),
+        ]
+    )
+    # データセットの読み込み
+    sample_dataset = datasets.ImageFolder(data_dir, transform=transform)
+
+    # データ拡張
+    for _ in range(num_augmentations):
+        sample_dataset += datasets.ImageFolder(data_dir, transform=transform_aug)
+
+    # データローダーの生成
+    data_loader = DataLoader(sample_dataset, batch_size=32, shuffle=True)
+
+    # モデルのファインチューニング
+    model.fit_s(
+        criterion=criterion,
+        optimizer=optimizer,
+        num_epochs=num_epochs,
+        learning_rate=learning_rate,
+        data_loader=data_loader,
+        device=model.device,
+        progress_bar=progress_bar,
+    )
+
+    return model

utils.py

@@ -1,13 +1,12 @@
 import os
-import tempfile
+from pathlib import Path
+
 from PIL import Image
 
 
-def initialize_data_dir() -> str:
-    if "data_dir" not in os.environ:
-        data_dir = tempfile.mkdtemp()
-        os.environ["data_dir"] = data_dir
-    return os.environ["data_dir"]
+def initialize_data_dir(data_dir: str) -> str:
+    Path(data_dir).mkdir(parents=True, exist_ok=True)
+    return data_dir
 
 
 def save_image(img, char, idx, data_dir):