chore: use bf16 with cuda

app.py

@@ -28,6 +28,7 @@ DEVICE = (
     if torch.backends.mps.is_available()
     else torch.device("cpu")
 )
+DTYPE = torch.bfloat16 if DEVICE.type in ["cuda"] else torch.float16
 MAX_TOKEN_LENGTH = 32
 
 model_map: dict[str, JiTModel] = {}  # {model_path: model}
@@ -64,6 +65,7 @@ def load_model(
     label2id_path: str,
     config_path: str,
     device: torch.device,
+    dtype: torch.dtype = DTYPE,
 ) -> tuple[JiTModel, dict]:
     """モデルを読み込む"""
 
@@ -84,7 +86,7 @@ def load_model(
     )
     model.eval()
     model.requires_grad_(False)
-    model.to(device=device)
+    model.to(device=device, dtype=dtype)
     model_map[model_path] = model  # cache
 
     label2id = load_label2id(get_file_path(MODEL_REPO, label2id_path))
@@ -113,9 +115,10 @@ def generate_images(
         label2id_path=label2id_path,
         config_path=config_path,
         device=DEVICE,
+        dtype=DTYPE,
     )
 
-    with torch.inference_mode():
+    with torch.inference_mode(), torch.autocast(device_type=DEVICE.type, dtype=DTYPE):
         images = model.generate(
             prompt=[prompt] * batch_size,
             negative_prompt=negative_prompt,
@@ -127,7 +130,7 @@ def generate_images(
             cfg_time_range=[0.1, 1.0],
             seed=seed if seed >= 0 else None,
             device=DEVICE,
-            execution_dtype=model.config.torch_dtype,
+            execution_dtype=DTYPE,
         )
 
     return images
@@ -271,6 +274,7 @@ if __name__ == "__main__":
         label2id_path=LABEL2ID_PATH,
         config_path=CONFIG_PATH,
         device=DEVICE,
+        dtype=DTYPE,
     )
 
     demo().launch()