From c84646dcd3ea295e6971b9609fc7bbcbcf3088bd Mon Sep 17 00:00:00 2001 From: hodanov <1031hoda@gmail.com> Date: Mon, 4 Nov 2024 12:20:04 +0900 Subject: [PATCH 1/2] Modify some instance variables to private. --- app/stable_diffusion_xl.py | 52 ++++++++++++++++++++------------------ 1 file changed, 28 insertions(+), 24 deletions(-) diff --git a/app/stable_diffusion_xl.py b/app/stable_diffusion_xl.py index e0df6cd..c113569 100644 --- a/app/stable_diffusion_xl.py +++ b/app/stable_diffusion_xl.py @@ -18,40 +18,40 @@ class SDXLTxt2Img: """ @enter() - def _setup(self): + def setup(self) -> None: import diffusers import torch import yaml config = {} - with open("/config.yml", "r") as file: + with open("/config.yml") 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.") + 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.") + print(f"The directory '{self.__cache_path}' does not exist.") - self.pipe = diffusers.StableDiffusionXLPipeline.from_pretrained( - self.cache_path, + self.__pipe = diffusers.StableDiffusionXLPipeline.from_pretrained( + self.__cache_path, torch_dtype=torch.float16, use_safetensors=True, ) - self.refiner = diffusers.StableDiffusionXLImg2ImgPipeline.from_pretrained( - self.cache_path, + self.__refiner = diffusers.StableDiffusionXLImg2ImgPipeline.from_pretrained( + self.__cache_path, torch_dtype=torch.float16, use_safetensors=True, ) - def _count_token(self, p: str, n: str) -> int: + def __count_token(self, p: str, n: str) -> int: """ Count the number of tokens in the prompt and negative prompt. """ from transformers import CLIPTokenizer tokenizer = CLIPTokenizer.from_pretrained( - self.cache_path, + self.__cache_path, subfolder="tokenizer", ) token_size_p = len(tokenizer.tokenize(p)) @@ -72,49 +72,53 @@ class SDXLTxt2Img: @method() def run_inference( self, + *, prompt: str, n_prompt: str, height: int = 1024, width: int = 1024, steps: int = 30, seed: int = 1, - use_upscaler: bool = False, output_format: str = "png", + use_upscaler: bool = False, ) -> list[bytes]: """ Runs the Stable Diffusion pipeline on the given prompt and outputs images. """ - import pillow_avif # noqa + import pillow_avif # noqa: F401 import torch + max_embeddings_multiples = self.__count_token(p=prompt, n=n_prompt) generator = torch.Generator("cuda").manual_seed(seed) - self.pipe.to("cuda") - self.pipe.enable_vae_tiling() - self.pipe.enable_xformers_memory_efficient_attention() - generated_image = self.pipe( + self.__pipe.to("cuda") + self.__pipe.enable_vae_tiling() + self.__pipe.enable_xformers_memory_efficient_attention() + generated_image = self.__pipe( prompt=prompt, negative_prompt=n_prompt, guidance_scale=7, height=height, width=width, generator=generator, + max_embeddings_multiples=max_embeddings_multiples, num_inference_steps=steps, ).images[0] generated_images = [generated_image] if use_upscaler: - self.refiner.to("cuda") - self.refiner.enable_vae_tiling() - self.refiner.enable_xformers_memory_efficient_attention() - base_image = self._double_image_size(generated_image) - image = self.refiner( + self.__refiner.to("cuda") + self.__refiner.enable_vae_tiling() + self.__refiner.enable_xformers_memory_efficient_attention() + base_image = self.__double_image_size(generated_image) + image = self.__refiner( prompt=prompt, negative_prompt=n_prompt, num_inference_steps=steps, strength=0.3, guidance_scale=7.5, generator=generator, + max_embeddings_multiples=max_embeddings_multiples, image=base_image, ).images[0] generated_images.append(image) @@ -127,7 +131,7 @@ class SDXLTxt2Img: return image_output - def _double_image_size(self, image: PIL.Image.Image) -> PIL.Image.Image: + def __double_image_size(self, image: PIL.Image.Image) -> PIL.Image.Image: image = image.convert("RGB") width, height = image.size return image.resize((width * 2, height * 2), resample=PIL.Image.LANCZOS) 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 2/2] 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,