Separate setup.py.

Makefile

@@ -1,5 +1,5 @@
 deploy:
-	modal deploy ./setup_files/setup.py
+	cd ./setup_files && modal deploy main.py
 
 # `--upscaler` is a name of upscaler you want to use.
 # You can use upscalers the below:

setup_files/main.py

@@ -0,0 +1,289 @@
+from __future__ import annotations
+
+import io
+import os
+
+import diffusers
+import PIL.Image
+import torch
+from modal import Secret, method
+from modal.cls import ClsMixin
+
+from setup import (BASE_CACHE_PATH, BASE_CACHE_PATH_CONTROLNET,
+                   BASE_CACHE_PATH_LORA, BASE_CACHE_PATH_TEXTUAL_INVERSION,
+                   stub)
+
+
+@stub.cls(
+    gpu="A10G",
+    secrets=[Secret.from_dotenv(__file__)],
+)
+class StableDiffusion(ClsMixin):
+    """
+    A class that wraps the Stable Diffusion pipeline and scheduler.
+    """
+
+    def __enter__(self):
+        import yaml
+
+        config = {}
+        with open("/config.yml", "r") 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.")
+
+        torch.cuda.memory._set_allocator_settings("max_split_size_mb:256")
+
+        self.pipe = diffusers.StableDiffusionPipeline.from_pretrained(
+            self.cache_path,
+            custom_pipeline="lpw_stable_diffusion",
+            torch_dtype=torch.float16,
+        )
+
+        # TODO: Add support for other schedulers.
+        self.pipe.scheduler = diffusers.EulerAncestralDiscreteScheduler.from_pretrained(
+            # self.pipe.scheduler = diffusers.DPMSolverMultistepScheduler.from_pretrained(
+            self.cache_path,
+            subfolder="scheduler",
+        )
+
+        vae = config.get("vae")
+        if vae is not None:
+            self.pipe.vae = diffusers.AutoencoderKL.from_pretrained(
+                self.cache_path,
+                subfolder="vae",
+            )
+        self.pipe.to("cuda")
+
+        loras = config.get("loras")
+        if loras is not None:
+            for lora in loras:
+                path = os.path.join(BASE_CACHE_PATH_LORA, lora["name"])
+                if os.path.exists(path):
+                    print(f"The directory '{path}' exists.")
+                else:
+                    print(f"The directory '{path}' does not exist. Need to execute 'modal deploy' first.")
+                self.pipe.load_lora_weights(".", weight_name=path)
+
+        textual_inversions = config.get("textual_inversions")
+        if textual_inversions is not None:
+            for textual_inversion in textual_inversions:
+                path = os.path.join(BASE_CACHE_PATH_TEXTUAL_INVERSION, textual_inversion["name"])
+                if os.path.exists(path):
+                    print(f"The directory '{path}' exists.")
+                else:
+                    print(f"The directory '{path}' does not exist. Need to execute 'modal deploy' first.")
+                self.pipe.load_textual_inversion(path)
+
+        self.pipe.enable_xformers_memory_efficient_attention()
+
+        # TODO: Repair the controlnet loading.
+        controlnets = config.get("controlnets")
+        if controlnets is not None:
+            for controlnet in controlnets:
+                path = os.path.join(BASE_CACHE_PATH_CONTROLNET, controlnet["name"])
+                controlnet = diffusers.ControlNetModel.from_pretrained(path, torch_dtype=torch.float16)
+                self.controlnet_pipe = diffusers.StableDiffusionControlNetPipeline.from_pretrained(
+                    self.cache_path,
+                    controlnet=controlnet,
+                    custom_pipeline="lpw_stable_diffusion",
+                    scheduler=self.pipe.scheduler,
+                    vae=self.pipe.vae,
+                    torch_dtype=torch.float16,
+                )
+                self.controlnet_pipe.to("cuda")
+                self.controlnet_pipe.enable_xformers_memory_efficient_attention()
+
+    @method()
+    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,
+            subfolder="tokenizer",
+        )
+        token_size_p = len(tokenizer.tokenize(p))
+        token_size_n = len(tokenizer.tokenize(n))
+        token_size = token_size_p
+        if token_size_p <= token_size_n:
+            token_size = token_size_n
+
+        max_embeddings_multiples = 1
+        max_length = tokenizer.model_max_length - 2
+        if token_size > max_length:
+            max_embeddings_multiples = token_size // max_length + 1
+
+        print(f"token_size: {token_size}, max_embeddings_multiples: {max_embeddings_multiples}")
+
+        return max_embeddings_multiples
+
+    @method()
+    def run_inference(
+        self,
+        prompt: str,
+        n_prompt: str,
+        height: int = 512,
+        width: int = 512,
+        samples: int = 1,
+        batch_size: int = 1,
+        steps: int = 30,
+        seed: int = 1,
+        upscaler: str = "",
+        use_face_enhancer: bool = False,
+        fix_by_controlnet_tile: bool = False,
+    ) -> list[bytes]:
+        """
+        Runs the Stable Diffusion pipeline on the given prompt and outputs images.
+        """
+
+        max_embeddings_multiples = self.count_token(p=prompt, n=n_prompt)
+        generator = torch.Generator("cuda").manual_seed(seed)
+        with torch.inference_mode():
+            with torch.autocast("cuda"):
+                generated_images = self.pipe.text2img(
+                    prompt * batch_size,
+                    negative_prompt=n_prompt * batch_size,
+                    height=height,
+                    width=width,
+                    num_inference_steps=steps,
+                    guidance_scale=7.5,
+                    max_embeddings_multiples=max_embeddings_multiples,
+                    generator=generator,
+                ).images
+
+        base_images = generated_images
+
+        """
+        Fix the generated images by the control_v11f1e_sd15_tile when `fix_by_controlnet_tile` is `True`.
+        https://huggingface.co/lllyasviel/control_v11f1e_sd15_tile
+        """
+        if fix_by_controlnet_tile:
+            for image in base_images:
+                image = self.resize_image(image=image, scale_factor=2)
+                with torch.inference_mode():
+                    with torch.autocast("cuda"):
+                        fixed_by_controlnet = self.controlnet_pipe(
+                            prompt=prompt * batch_size,
+                            negative_prompt=n_prompt * batch_size,
+                            num_inference_steps=steps,
+                            strength=0.3,
+                            guidance_scale=7.5,
+                            max_embeddings_multiples=max_embeddings_multiples,
+                            generator=generator,
+                            image=image,
+                        ).images
+            generated_images.extend(fixed_by_controlnet)
+            base_images = fixed_by_controlnet
+
+        if upscaler != "":
+            upscaled = self.upscale(
+                base_images=base_images,
+                half_precision=False,
+                tile=700,
+                upscaler=upscaler,
+                use_face_enhancer=use_face_enhancer,
+            )
+            generated_images.extend(upscaled)
+
+        image_output = []
+        for image in generated_images:
+            with io.BytesIO() as buf:
+                image.save(buf, format="PNG")
+                image_output.append(buf.getvalue())
+
+        return image_output
+
+    @method()
+    def resize_image(self, image: PIL.Image.Image, scale_factor: int) -> PIL.Image.Image:
+        image = image.convert("RGB")
+        width, height = image.size
+        img = image.resize((width * scale_factor, height * scale_factor), resample=PIL.Image.LANCZOS)
+        return img
+
+    @method()
+    def upscale(
+        self,
+        base_images: list[PIL.Image],
+        half_precision: bool = False,
+        tile: int = 0,
+        tile_pad: int = 10,
+        pre_pad: int = 0,
+        upscaler: str = "",
+        use_face_enhancer: bool = False,
+    ) -> list[PIL.Image]:
+        """
+        Upscale the generated images by the upscaler when `upscaler` is selected.
+        The upscaler can be selected from the following list:
+        - `RealESRGAN_x4plus`
+        - `RealESRNet_x4plus`
+        - `RealESRGAN_x4plus_anime_6B`
+        - `RealESRGAN_x2plus`
+        https://github.com/xinntao/Real-ESRGAN
+        """
+        import numpy
+        from basicsr.archs.rrdbnet_arch import RRDBNet
+        from gfpgan import GFPGANer
+        from realesrgan import RealESRGANer
+        from tqdm import tqdm
+
+        model_name = upscaler
+        if model_name == "RealESRGAN_x4plus":
+            upscale_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+            netscale = 4
+        elif model_name == "RealESRNet_x4plus":
+            upscale_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+            netscale = 4
+        elif model_name == "RealESRGAN_x4plus_anime_6B":
+            upscale_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
+            netscale = 4
+        elif model_name == "RealESRGAN_x2plus":
+            upscale_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
+            netscale = 2
+        else:
+            raise NotImplementedError("Model name not supported")
+
+        upsampler = RealESRGANer(
+            scale=netscale,
+            model_path=os.path.join(BASE_CACHE_PATH, "esrgan", f"{model_name}.pth"),
+            dni_weight=None,
+            model=upscale_model,
+            tile=tile,
+            tile_pad=tile_pad,
+            pre_pad=pre_pad,
+            half=half_precision,
+            gpu_id=None,
+        )
+
+        if use_face_enhancer:
+            face_enhancer = GFPGANer(
+                model_path=os.path.join(BASE_CACHE_PATH, "esrgan", "GFPGANv1.3.pth"),
+                upscale=netscale,
+                arch="clean",
+                channel_multiplier=2,
+                bg_upsampler=upsampler,
+            )
+
+        upscaled_imgs = []
+        with tqdm(total=len(base_images)) as progress_bar:
+            for img in base_images:
+                img = numpy.array(img)
+                if use_face_enhancer:
+                    _, _, enhance_result = face_enhancer.enhance(
+                        img,
+                        has_aligned=False,
+                        only_center_face=False,
+                        paste_back=True,
+                    )
+                else:
+                    enhance_result, _ = upsampler.enhance(img)
+
+                upscaled_imgs.append(PIL.Image.fromarray(enhance_result))
+                progress_bar.update(1)
+
+        return upscaled_imgs

setup_files/setup.py

@@ -1,13 +1,9 @@
 from __future__ import annotations
 
-import io
 import os
-from urllib.request import Request, urlopen
 
 import diffusers
-import yaml
+from modal import Image, Mount, Secret, Stub
-from modal import Image, Mount, Secret, Stub, method
-from modal.cls import ClsMixin
 
 BASE_CACHE_PATH = "/vol/cache"
 BASE_CACHE_PATH_LORA = "/vol/cache/lora"
@@ -19,6 +15,8 @@ def download_file(url, file_name, file_path):
     """
     Download files.
     """
+    from urllib.request import Request, urlopen
+
     req = Request(url, headers={"User-Agent": "Mozilla/5.0"})
     downloaded = urlopen(req).read()
     dir_names = os.path.join(file_path, file_name)
@@ -70,6 +68,8 @@ def build_image():
     """
     Build the Docker image.
     """
+    import yaml
+
     token = os.environ["HUGGING_FACE_TOKEN"]
     config = {}
     with open("/config.yml", "r") as file:
@@ -109,302 +109,16 @@ def build_image():
 
 stub = Stub("stable-diffusion-cli")
 base_stub = Image.from_dockerfile(
-    path="./setup_files/Dockerfile",
+    path="Dockerfile",
-    context_mount=Mount.from_local_file("./setup_files/requirements.txt"),
+    context_mount=Mount.from_local_file("requirements.txt"),
 )
 stub.image = base_stub.extend(
     dockerfile_commands=[
         "FROM base",
-        "COPY ./config.yml /",
+        "COPY config.yml /",
     ],
-    context_mount=Mount.from_local_file("./setup_files/config.yml"),
+    context_mount=Mount.from_local_file("config.yml"),
 ).run_function(
     build_image,
     secrets=[Secret.from_dotenv(__file__)],
 )
-
-
-@stub.cls(
-    gpu="A10G",
-    secrets=[Secret.from_dotenv(__file__)],
-)
-class StableDiffusion(ClsMixin):
-    """
-    A class that wraps the Stable Diffusion pipeline and scheduler.
-    """
-
-    def __enter__(self):
-        import torch
-
-        config = {}
-        with open("/config.yml", "r") 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.")
-
-        torch.cuda.memory._set_allocator_settings("max_split_size_mb:256")
-
-        self.pipe = diffusers.StableDiffusionPipeline.from_pretrained(
-            self.cache_path,
-            custom_pipeline="lpw_stable_diffusion",
-            torch_dtype=torch.float16,
-        )
-
-        # TODO: Add support for other schedulers.
-        self.pipe.scheduler = diffusers.EulerAncestralDiscreteScheduler.from_pretrained(
-            # self.pipe.scheduler = diffusers.DPMSolverMultistepScheduler.from_pretrained(
-            self.cache_path,
-            subfolder="scheduler",
-        )
-
-        vae = config.get("vae")
-        if vae is not None:
-            self.pipe.vae = diffusers.AutoencoderKL.from_pretrained(
-                self.cache_path,
-                subfolder="vae",
-            )
-        self.pipe.to("cuda")
-
-        loras = config.get("loras")
-        if loras is not None:
-            for lora in loras:
-                path = os.path.join(BASE_CACHE_PATH_LORA, lora["name"])
-                if os.path.exists(path):
-                    print(f"The directory '{path}' exists.")
-                else:
-                    print(f"The directory '{path}' does not exist. Download it...")
-                    download_file(lora["download_url"], lora["name"], BASE_CACHE_PATH_LORA)
-                self.pipe.load_lora_weights(".", weight_name=path)
-
-        textual_inversions = config.get("textual_inversions")
-        if textual_inversions is not None:
-            for textual_inversion in textual_inversions:
-                path = os.path.join(BASE_CACHE_PATH_TEXTUAL_INVERSION, textual_inversion["name"])
-                if os.path.exists(path):
-                    print(f"The directory '{path}' exists.")
-                else:
-                    print(f"The directory '{path}' does not exist. Download it...")
-                    download_file(
-                        textual_inversion["download_url"],
-                        textual_inversion["name"],
-                        BASE_CACHE_PATH_TEXTUAL_INVERSION,
-                    )
-                self.pipe.load_textual_inversion(path)
-
-        self.pipe.enable_xformers_memory_efficient_attention()
-
-        # TODO: Repair the controlnet loading.
-        controlnets = config.get("controlnets")
-        if controlnets is not None:
-            for controlnet in controlnets:
-                path = os.path.join(BASE_CACHE_PATH_CONTROLNET, controlnet["name"])
-                controlnet = diffusers.ControlNetModel.from_pretrained(path, torch_dtype=torch.float16)
-                self.controlnet_pipe = diffusers.StableDiffusionControlNetPipeline.from_pretrained(
-                    self.cache_path,
-                    controlnet=controlnet,
-                    custom_pipeline="lpw_stable_diffusion",
-                    scheduler=self.pipe.scheduler,
-                    vae=self.pipe.vae,
-                    torch_dtype=torch.float16,
-                )
-                self.controlnet_pipe.to("cuda")
-                self.controlnet_pipe.enable_xformers_memory_efficient_attention()
-
-    @method()
-    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,
-            subfolder="tokenizer",
-        )
-        token_size_p = len(tokenizer.tokenize(p))
-        token_size_n = len(tokenizer.tokenize(n))
-        token_size = token_size_p
-        if token_size_p <= token_size_n:
-            token_size = token_size_n
-
-        max_embeddings_multiples = 1
-        max_length = tokenizer.model_max_length - 2
-        if token_size > max_length:
-            max_embeddings_multiples = token_size // max_length + 1
-
-        print(f"token_size: {token_size}, max_embeddings_multiples: {max_embeddings_multiples}")
-
-        return max_embeddings_multiples
-
-    @method()
-    def run_inference(
-        self,
-        prompt: str,
-        n_prompt: str,
-        height: int = 512,
-        width: int = 512,
-        samples: int = 1,
-        batch_size: int = 1,
-        steps: int = 30,
-        seed: int = 1,
-        upscaler: str = "",
-        use_face_enhancer: bool = False,
-        fix_by_controlnet_tile: bool = False,
-    ) -> list[bytes]:
-        """
-        Runs the Stable Diffusion pipeline on the given prompt and outputs images.
-        """
-        import torch
-
-        max_embeddings_multiples = self.count_token(p=prompt, n=n_prompt)
-        generator = torch.Generator("cuda").manual_seed(seed)
-        with torch.inference_mode():
-            with torch.autocast("cuda"):
-                generated_images = self.pipe.text2img(
-                    prompt * batch_size,
-                    negative_prompt=n_prompt * batch_size,
-                    height=height,
-                    width=width,
-                    num_inference_steps=steps,
-                    guidance_scale=7.5,
-                    max_embeddings_multiples=max_embeddings_multiples,
-                    generator=generator,
-                ).images
-
-        base_images = generated_images
-
-        """
-        Fix the generated images by the control_v11f1e_sd15_tile when `fix_by_controlnet_tile` is `True`.
-        https://huggingface.co/lllyasviel/control_v11f1e_sd15_tile
-        """
-        if fix_by_controlnet_tile:
-            for image in base_images:
-                image = self.resize_image(image=image, scale_factor=2)
-                with torch.inference_mode():
-                    with torch.autocast("cuda"):
-                        fixed_by_controlnet = self.controlnet_pipe(
-                            prompt=prompt * batch_size,
-                            negative_prompt=n_prompt * batch_size,
-                            num_inference_steps=steps,
-                            strength=0.3,
-                            guidance_scale=7.5,
-                            max_embeddings_multiples=max_embeddings_multiples,
-                            generator=generator,
-                            image=image,
-                        ).images
-            generated_images.extend(fixed_by_controlnet)
-            base_images = fixed_by_controlnet
-
-        if upscaler != "":
-            upscaled = self.upscale(
-                base_images=base_images,
-                half_precision=False,
-                tile=700,
-                upscaler=upscaler,
-                use_face_enhancer=use_face_enhancer,
-            )
-            generated_images.extend(upscaled)
-
-        image_output = []
-        for image in generated_images:
-            with io.BytesIO() as buf:
-                image.save(buf, format="PNG")
-                image_output.append(buf.getvalue())
-
-        return image_output
-
-    @method()
-    def resize_image(self, image: Image.Image, scale_factor: int) -> Image.Image:
-        from PIL import Image
-
-        image = image.convert("RGB")
-        width, height = image.size
-        img = image.resize((width * scale_factor, height * scale_factor), resample=Image.LANCZOS)
-        return img
-
-    @method()
-    def upscale(
-        self,
-        base_images: list[Image.Image],
-        half_precision: bool = False,
-        tile: int = 0,
-        tile_pad: int = 10,
-        pre_pad: int = 0,
-        upscaler: str = "",
-        use_face_enhancer: bool = False,
-    ) -> list[Image.Image]:
-        """
-        Upscale the generated images by the upscaler when `upscaler` is selected.
-        The upscaler can be selected from the following list:
-        - `RealESRGAN_x4plus`
-        - `RealESRNet_x4plus`
-        - `RealESRGAN_x4plus_anime_6B`
-        - `RealESRGAN_x2plus`
-        https://github.com/xinntao/Real-ESRGAN
-        """
-        import numpy
-        from basicsr.archs.rrdbnet_arch import RRDBNet
-        from PIL import Image
-        from realesrgan import RealESRGANer
-        from tqdm import tqdm
-
-        model_name = upscaler
-        if model_name == "RealESRGAN_x4plus":
-            upscale_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
-            netscale = 4
-        elif model_name == "RealESRNet_x4plus":
-            upscale_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
-            netscale = 4
-        elif model_name == "RealESRGAN_x4plus_anime_6B":
-            upscale_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
-            netscale = 4
-        elif model_name == "RealESRGAN_x2plus":
-            upscale_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
-            netscale = 2
-        else:
-            raise NotImplementedError("Model name not supported")
-
-        upsampler = RealESRGANer(
-            scale=netscale,
-            model_path=os.path.join(BASE_CACHE_PATH, "esrgan", f"{model_name}.pth"),
-            dni_weight=None,
-            model=upscale_model,
-            tile=tile,
-            tile_pad=tile_pad,
-            pre_pad=pre_pad,
-            half=half_precision,
-            gpu_id=None,
-        )
-
-        from gfpgan import GFPGANer
-
-        if use_face_enhancer:
-            face_enhancer = GFPGANer(
-                model_path=os.path.join(BASE_CACHE_PATH, "esrgan", "GFPGANv1.3.pth"),
-                upscale=netscale,
-                arch="clean",
-                channel_multiplier=2,
-                bg_upsampler=upsampler,
-            )
-
-        upscaled_imgs = []
-        with tqdm(total=len(base_images)) as progress_bar:
-            for img in base_images:
-                img = numpy.array(img)
-                if use_face_enhancer:
-                    _, _, enhance_result = face_enhancer.enhance(
-                        img,
-                        has_aligned=False,
-                        only_center_face=False,
-                        paste_back=True,
-                    )
-                else:
-                    enhance_result, _ = upsampler.enhance(img)
-
-                upscaled_imgs.append(Image.fromarray(enhance_result))
-                progress_bar.update(1)
-
-        return upscaled_imgs