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

Fixed a bug when using haggingface token.

Browse Source
This commit is contained in:
hodanov 2024-11-04 11:43:49 +09:00
parent 22b8ab721f
commit afa0053ce8
1 changed files with 12 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
app/setup.py
View File

@ -1,8 +1,8 @@
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 modal import App, Image, Mount, Secret from modal import App, Image, Mount, Secret
BASE_CACHE_PATH = "/vol/cache" BASE_CACHE_PATH = "/vol/cache"
@ -30,10 +30,13 @@ class StableDiffusionCLISetupSDXL(StableDiffusionCLISetupInterface):
self.__model_name: str = config["model"]["name"] self.__model_name: str = config["model"]["name"]
self.__model_url: str = config["model"]["url"] self.__model_url: str = config["model"]["url"]
if token != "":
login(token)
self.__token: str = token self.__token: str = token
def download_model(self) -> None: def download_model(self) -> None:
cache_path = Path(BASE_CACHE_PATH, self.__model_name) cache_path = os.path.join(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,
@ -54,10 +57,13 @@ class StableDiffusionCLISetupSD15(StableDiffusionCLISetupInterface):
self.__model_name: str = config["model"]["name"] self.__model_name: str = config["model"]["name"]
self.__model_url: str = config["model"]["url"] self.__model_url: str = config["model"]["url"]
if token != "":
login(token)
self.__token: str = token self.__token: str = token
def download_model(self) -> None: def download_model(self) -> None:
cache_path = Path(BASE_CACHE_PATH, self.__model_name) cache_path = os.path.join(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,
@ -111,7 +117,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 = Path(BASE_CACHE_PATH, name) cache_path = os.path.join(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,
@ -120,7 +126,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 = Path(BASE_CACHE_PATH, name) cache_path = os.path.join(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,
@ -136,7 +142,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 = Path(file_path, file_name) dir_names = os.path.join(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)