From 871be963218577a738af951fc67457b835356d97 Mon Sep 17 00:00:00 2001
From: hodanov <1031hoda@gmail.com>
Date: Mon, 4 Nov 2024 14:10:32 +0900
Subject: [PATCH] Refactor app/stable_diffusion_xl.py.

---
 app/setup.py               | 11 ++++++-----
 app/stable_diffusion_xl.py | 13 ++++++-------
 2 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/app/setup.py b/app/setup.py
index 5e0939e..8718a4b 100644
--- a/app/setup.py
+++ b/app/setup.py
@@ -1,5 +1,6 @@
 import os
 from abc import ABC, abstractmethod
+from pathlib import Path
 
 import diffusers
 from huggingface_hub import login
@@ -36,7 +37,7 @@ class StableDiffusionCLISetupSDXL(StableDiffusionCLISetupInterface):
         self.__token: str = token
 
     def download_model(self) -> None:
-        cache_path = os.path.join(BASE_CACHE_PATH, self.__model_name)
+        cache_path = Path(BASE_CACHE_PATH) / self.__model_name
         pipe = diffusers.StableDiffusionXLPipeline.from_single_file(
             pretrained_model_link_or_path=self.__model_url,
             use_auth_token=self.__token,
@@ -63,7 +64,7 @@ class StableDiffusionCLISetupSD15(StableDiffusionCLISetupInterface):
         self.__token: str = token
 
     def download_model(self) -> None:
-        cache_path = os.path.join(BASE_CACHE_PATH, self.__model_name)
+        cache_path = Path(BASE_CACHE_PATH) / self.__model_name
         pipe = diffusers.StableDiffusionPipeline.from_single_file(
             pretrained_model_link_or_path=self.__model_url,
             token=self.__token,
@@ -117,7 +118,7 @@ class CommonSetup:
                 )
 
     def __download_vae(self, name: str, model_url: str, token: str) -> None:
-        cache_path = os.path.join(BASE_CACHE_PATH, name)
+        cache_path = Path(BASE_CACHE_PATH, name)
         vae = diffusers.AutoencoderKL.from_single_file(
             pretrained_model_link_or_path=model_url,
             use_auth_token=token,
@@ -126,7 +127,7 @@ class CommonSetup:
         vae.save_pretrained(cache_path, safe_serialization=True)
 
     def __download_controlnet(self, name: str, repo_id: str, token: str) -> None:
-        cache_path = os.path.join(BASE_CACHE_PATH_CONTROLNET, name)
+        cache_path = Path(BASE_CACHE_PATH_CONTROLNET) / name
         controlnet = diffusers.ControlNetModel.from_pretrained(
             repo_id,
             use_auth_token=token,
@@ -142,7 +143,7 @@ class CommonSetup:
 
         req = Request(url, headers={"User-Agent": "Mozilla/5.0"})
         downloaded = urlopen(req).read()
-        dir_names = os.path.join(file_path, file_name)
+        dir_names = Path(file_path) / file_name
         os.makedirs(os.path.dirname(dir_names), exist_ok=True)
         with open(dir_names, mode="wb") as f:
             f.write(downloaded)
diff --git a/app/stable_diffusion_xl.py b/app/stable_diffusion_xl.py
index c113569..e23cd26 100644
--- a/app/stable_diffusion_xl.py
+++ b/app/stable_diffusion_xl.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import io
-import os
+from pathlib import Path
 
 import PIL.Image
 from modal import Secret, enter, method
@@ -24,13 +24,12 @@ class SDXLTxt2Img:
         import yaml
 
         config = {}
-        with open("/config.yml") as file:
+        with Path("/config.yml").open() as file:
             config = yaml.safe_load(file)
-        self.__cache_path = os.path.join(BASE_CACHE_PATH, config["model"]["name"])
-        if os.path.exists(self.__cache_path):
-            print(f"The directory '{self.__cache_path}' exists.")
-        else:
-            print(f"The directory '{self.__cache_path}' does not exist.")
+        self.__cache_path = Path(BASE_CACHE_PATH) / config["model"]["name"]
+        if not Path.exists(self.__cache_path):
+            msg = f"The directory '{self.__cache_path}' does not exist."
+            raise ValueError(msg)
 
         self.__pipe = diffusers.StableDiffusionXLPipeline.from_pretrained(
             self.__cache_path,