chunzhi/stable-diffusion-cli-on-modal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor app/stable_diffusion_xl.py.

Browse Source
This commit is contained in:
hodanov 2024-11-04 14:10:32 +09:00
parent c84646dcd3
commit 871be96321
2 changed files with 12 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
app/setup.py
View File

@ -1,5 +1,6 @@
import os import os
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from pathlib import Path
import diffusers import diffusers
from huggingface_hub import login from huggingface_hub import login
@ -36,7 +37,7 @@ class StableDiffusionCLISetupSDXL(StableDiffusionCLISetupInterface):
self.__token: str = token self.__token: str = token
def download_model(self) -> None: 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( pipe = diffusers.StableDiffusionXLPipeline.from_single_file(
pretrained_model_link_or_path=self.__model_url, pretrained_model_link_or_path=self.__model_url,
use_auth_token=self.__token, use_auth_token=self.__token,
@ -63,7 +64,7 @@ class StableDiffusionCLISetupSD15(StableDiffusionCLISetupInterface):
self.__token: str = token self.__token: str = token
def download_model(self) -> None: 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( pipe = diffusers.StableDiffusionPipeline.from_single_file(
pretrained_model_link_or_path=self.__model_url, pretrained_model_link_or_path=self.__model_url,
token=self.__token, token=self.__token,
@ -117,7 +118,7 @@ class CommonSetup:
) )
def __download_vae(self, name: str, model_url: str, token: str) -> None: 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( vae = diffusers.AutoencoderKL.from_single_file(
pretrained_model_link_or_path=model_url, pretrained_model_link_or_path=model_url,
use_auth_token=token, use_auth_token=token,
@ -126,7 +127,7 @@ class CommonSetup:
vae.save_pretrained(cache_path, safe_serialization=True) vae.save_pretrained(cache_path, safe_serialization=True)
def __download_controlnet(self, name: str, repo_id: str, token: str) -> None: 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( controlnet = diffusers.ControlNetModel.from_pretrained(
repo_id, repo_id,
use_auth_token=token, use_auth_token=token,
@ -142,7 +143,7 @@ class CommonSetup:
req = Request(url, headers={"User-Agent": "Mozilla/5.0"}) req = Request(url, headers={"User-Agent": "Mozilla/5.0"})
downloaded = urlopen(req).read() 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) os.makedirs(os.path.dirname(dir_names), exist_ok=True)
with open(dir_names, mode="wb") as f: with open(dir_names, mode="wb") as f:
f.write(downloaded) f.write(downloaded)

13
app/stable_diffusion_xl.py
View File

@ -1,7 +1,7 @@
from __future__ import annotations from __future__ import annotations
import io import io
import os from pathlib import Path
import PIL.Image import PIL.Image
from modal import Secret, enter, method from modal import Secret, enter, method
@ -24,13 +24,12 @@ class SDXLTxt2Img:
import yaml import yaml
config = {} config = {}
with open("/config.yml") as file: with Path("/config.yml").open() as file:
config = yaml.safe_load(file) config = yaml.safe_load(file)
self.__cache_path = os.path.join(BASE_CACHE_PATH, config["model"]["name"]) self.__cache_path = Path(BASE_CACHE_PATH) / config["model"]["name"]
if os.path.exists(self.__cache_path): if not Path.exists(self.__cache_path):
print(f"The directory '{self.__cache_path}' exists.") msg = f"The directory '{self.__cache_path}' does not exist."
else: raise ValueError(msg)
print(f"The directory '{self.__cache_path}' does not exist.")
self.__pipe = diffusers.StableDiffusionXLPipeline.from_pretrained( self.__pipe = diffusers.StableDiffusionXLPipeline.from_pretrained(
self.__cache_path, self.__cache_path,