Implement SDXLTxt2Img.

Makefile

@@ -1,4 +1,4 @@
-deploy:
+app:
 	cd ./setup_files && modal deploy __main__.py
 
 # `--upscaler` is a name of upscaler you want to use.
@@ -7,8 +7,8 @@ deploy:
 #   - `RealESRNet_x4plus`
 #   - `RealESRGAN_x4plus_anime_6B`
 #   - `RealESRGAN_x2plus`
-run:
+img_by_sd15_txt2img:
-	cd ./sdcli && modal run txt2img.py \
+	cd ./sdcli && modal run sd15_txt2img.py \
 	--prompt "a photograph of an astronaut riding a horse" \
 	--n-prompt "" \
 	--height 512 \
@@ -17,4 +17,15 @@ run:
 	--steps 30 \
 	--upscaler "RealESRGAN_x2plus" \
 	--use-face-enhancer "False" \
-	--fix-by-controlnet-tile "True"
+	--fix-by-controlnet-tile "True" \
+	--output-format "avif"
+
+
+img_by_sdxl_txt2img:
+	cd ./sdcli && modal run sdxl_txt2img.py \
+	--prompt "A dog is running on the grass" \
+	--height 1024 \
+	--width 1024 \
+	--samples 1 \
+	--upscaler "RealESRGAN_x2plus" \
+	--output-format "avif"

sdcli/sd15_txt2img.py

@@ -4,7 +4,7 @@ import modal
 import util
 
 stub = modal.Stub("run-stable-diffusion-cli")
-stub.run_inference = modal.Function.from_name("stable-diffusion-cli", "Txt2Img.run_inference")
+stub.run_inference = modal.Function.from_name("stable-diffusion-cli", "SD15Txt2Img.run_inference")
 
 
 @stub.local_entrypoint()

sdcli/sdxl_txt2img.py

@@ -0,0 +1,51 @@
+import time
+
+import modal
+import util
+
+stub = modal.Stub("run-stable-diffusion-cli")
+stub.run_inference = modal.Function.from_name("stable-diffusion-cli", "SDXLTxt2Img.run_inference")
+
+
+@stub.local_entrypoint()
+def main(
+    prompt: str,
+    height: int = 1024,
+    width: int = 1024,
+    samples: int = 5,
+    seed: int = -1,
+    upscaler: str = "",
+    use_face_enhancer: str = "False",
+    output_format: str = "png",
+):
+    """
+    This function is the entrypoint for the Runway CLI.
+    The function pass the given prompt to StableDiffusion on Modal,
+    gets back a list of images and outputs images to local.
+    """
+    directory = util.make_directory()
+    seed_generated = seed
+    for i in range(samples):
+        if seed == -1:
+            seed_generated = util.generate_seed()
+        start_time = time.time()
+        images = stub.run_inference.remote(
+            prompt=prompt,
+            height=height,
+            width=width,
+            seed=seed_generated,
+            upscaler=upscaler,
+            use_face_enhancer=use_face_enhancer == "True",
+            output_format=output_format,
+        )
+        util.save_images(directory, images, seed_generated, i, output_format)
+        total_time = time.time() - start_time
+        print(f"Sample {i} took {total_time:.3f}s ({(total_time)/len(images):.3f}s / image).")
+
+    prompts: dict[str, int | str] = {
+        "prompt": prompt,
+        "height": height,
+        "width": width,
+        "samples": samples,
+    }
+    util.save_prompts(prompts)

setup_files/__main__.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
+import stable_diffusion_1_5
+import stable_diffusion_xl
 from setup import stub
-from stable_diffusion_1_5 import Txt2Img
 
 
 @stub.function(gpu="A10G")
 def main():
-    Txt2Img
+    stable_diffusion_1_5.SD15Txt2Img
+    stable_diffusion_xl.SDXLTxt2Img
 
 
 if __name__ == "__main__":

setup_files/setup.py

@@ -64,6 +64,26 @@ def download_model(name: str, model_url: str, token: str):
     pipe.save_pretrained(cache_path, safe_serialization=True)
 
 
+def download_model_sdxl(name: str, model_url: str, token: str):
+    """
+    Download a sdxl model.
+    """
+    cache_path = os.path.join(BASE_CACHE_PATH, name)
+    pipe = diffusers.StableDiffusionXLPipeline.from_single_file(
+        pretrained_model_link_or_path=model_url,
+        use_auth_token=token,
+        cache_dir=cache_path,
+    )
+    pipe.save_pretrained(cache_path, safe_serialization=True)
+
+    refiner_cache_path = cache_path + "-refiner"
+    refiner = diffusers.StableDiffusionXLImg2ImgPipeline.from_single_file(
+        "https://huggingface.co/stabilityai/stable-diffusion-xl-refiner-1.0/blob/main/sd_xl_refiner_1.0.safetensors",
+        cache_dir=refiner_cache_path,
+    )
+    refiner.save_pretrained(refiner_cache_path, safe_serialization=True)
+
+
 def build_image():
     """
     Build the Docker image.
@@ -76,8 +96,12 @@ def build_image():
         config = yaml.safe_load(file)
 
     model = config.get("model")
+    use_xl = config.get("use_xl")
     if model is not None:
-        download_model(name=model["name"], model_url=model["url"], token=token)
+        if use_xl is not None and use_xl:
+            download_model_sdxl(name=model["name"], model_url=model["url"], token=token)
+        else:
+            download_model(name=model["name"], model_url=model["url"], token=token)
 
     vae = config.get("vae")
     if vae is not None:

setup_files/stable_diffusion_1_5.py

@@ -18,7 +18,7 @@ from setup import (
     gpu="A10G",
     secrets=[Secret.from_dotenv(__file__)],
 )
-class Txt2Img:
+class SD15Txt2Img:
     """
     A class that wraps the Stable Diffusion pipeline and scheduler.
     """

setup_files/stable_diffusion_xl.py

@@ -0,0 +1,180 @@
+from __future__ import annotations
+
+import io
+import os
+
+import PIL.Image
+from modal import Secret, method
+from setup import BASE_CACHE_PATH, stub
+
+
+@stub.cls(
+    gpu="A10G",
+    secrets=[Secret.from_dotenv(__file__)],
+)
+class SDXLTxt2Img:
+    """
+    A class that wraps the Stable Diffusion pipeline and scheduler.
+    """
+
+    def __enter__(self):
+        import diffusers
+        import torch
+        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.")
+
+        self.pipe = diffusers.AutoPipelineForText2Image.from_pretrained(
+            self.cache_path,
+            torch_dtype=torch.float16,
+            use_safetensors=True,
+            variant="fp16",
+        )
+
+        self.refiner_cache_path = self.cache_path + "-refiner"
+        self.refiner = diffusers.StableDiffusionXLImg2ImgPipeline.from_pretrained(
+            self.refiner_cache_path,
+            torch_dtype=torch.float16,
+            use_safetensors=True,
+            variant="fp16",
+        )
+
+    @method()
+    def run_inference(
+        self,
+        prompt: str,
+        height: int = 1024,
+        width: int = 1024,
+        seed: int = 1,
+        upscaler: str = "",
+        use_face_enhancer: bool = False,
+        output_format: str = "png",
+    ) -> list[bytes]:
+        """
+        Runs the Stable Diffusion pipeline on the given prompt and outputs images.
+        """
+        import pillow_avif  # noqa
+        import torch
+
+        generator = torch.Generator("cuda").manual_seed(seed)
+        self.pipe.to("cuda")
+        generated_images = self.pipe(
+            prompt=prompt,
+            height=height,
+            width=width,
+            generator=generator,
+        ).images
+        base_images = generated_images
+
+        for image in base_images:
+            self.refiner.to("cuda")
+            refined_images = self.refiner(
+                prompt=prompt,
+                image=image,
+            ).images
+        generated_images.extend(refined_images)
+        base_images = refined_images
+
+        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=output_format)
+                image_output.append(buf.getvalue())
+
+        return image_output
+
+    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 = []
+        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))
+
+        return upscaled_imgs