Merge pull request #4 from hodanov/feature/add_sd_x2_latent_upscaler

Add Real-ESRGAN

Dockerfile

@@ -1,5 +1,10 @@
 FROM python:3.11.3-slim-bullseye
 COPY requirements.txt /
 RUN apt update \
-    && apt install -y wget git \
+    && apt install -y wget git libgl1-mesa-glx libglib2.0-0 \
-    && pip install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu117 --pre xformers
+    && pip install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu117 \
+    && mkdir -p /vol/cache/esrgan \
+    && wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P /vol/cache/esrgan \
+    && wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth -P /vol/cache/esrgan \
+    && wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth -P /vol/cache/esrgan \
+    && wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth -P /vol/cache/esrgan

Makefile

@@ -1,9 +1,9 @@
 run:
 	modal run sd_cli.py \
-	--prompt "a woman with bob hair" \
+	--prompt "A woman with bob hair" \
 	--n-prompt "" \
 	--height 768 \
 	--width 512 \
 	--samples 5 \
-	--steps 20 \
+	--steps 50 \
-	--upscaler "sd_x2_latent_upscaler"
+	--upscaler "RealESRGAN_x4plus_anime_6B"

requirements.txt

@@ -1,9 +1,17 @@
 accelerate
-scipy
+diffusers[torch]==0.16.1
-diffusers[torch]
+onnxruntime==1.15.0
-safetensors
+safetensors==0.3.1
 torch==2.0.1+cu117
+transformers==4.29.2
+xformers==0.0.20
+
+realesrgan
+basicsr>=1.4.2
+facexlib>=0.2.5
+gfpgan>=1.3.5
+numpy
+opencv-python
+Pillow
 torchvision
-torchmetrics
+tqdm
-omegaconf
-transformers

sd_cli.py

@@ -22,6 +22,13 @@ def download_models():
     model_repo_id = os.environ["MODEL_REPO_ID"]
     cache_path = os.path.join(BASE_CACHE_PATH, os.environ["MODEL_NAME"])
 
+    vae = diffusers.AutoencoderKL.from_pretrained(
+        "stabilityai/sd-vae-ft-mse",
+        use_auth_token=hugging_face_token,
+        cache_dir=cache_path,
+    )
+    vae.save_pretrained(cache_path, safe_serialization=True)
+
     scheduler = diffusers.EulerAncestralDiscreteScheduler.from_pretrained(
         model_repo_id,
         subfolder="scheduler",
@@ -68,6 +75,11 @@ class StableDiffusion:
 
         torch.backends.cuda.matmul.allow_tf32 = True
 
+        vae = diffusers.AutoencoderKL.from_pretrained(
+            cache_path,
+            subfolder="vae",
+        )
+
         scheduler = diffusers.EulerAncestralDiscreteScheduler.from_pretrained(
             cache_path,
             subfolder="scheduler",
@@ -76,21 +88,12 @@ class StableDiffusion:
         self.pipe = diffusers.StableDiffusionPipeline.from_pretrained(
             cache_path,
             scheduler=scheduler,
+            vae=vae,
             custom_pipeline="lpw_stable_diffusion",
+            torch_dtype=torch.float16,
         ).to("cuda")
         self.pipe.enable_xformers_memory_efficient_attention()
 
-        self.upscaler = diffusers.StableDiffusionLatentUpscalePipeline.from_pretrained(
-            "stabilityai/sd-x2-latent-upscaler", torch_dtype=torch.float16
-        ).to("cuda")
-        self.upscaler.enable_xformers_memory_efficient_attention()
-
-        # model_id = "stabilityai/stable-diffusion-x4-upscaler"
-        # self.upscaler = diffusers.StableDiffusionUpscalePipeline.from_pretrained(
-        #     , revision="fp16", torch_dtype=torch.float16
-        # ).to("cuda")
-        # self.upscaler.enable_xformers_memory_efficient_attention()
-
     @method()
     def run_inference(self, inputs: dict[str, int | str]) -> list[bytes]:
         """
@@ -100,7 +103,7 @@ class StableDiffusion:
 
         with torch.inference_mode():
             with torch.autocast("cuda"):
-                images = self.pipe(
+                base_images = self.pipe(
                     [inputs["prompt"]] * int(inputs["batch_size"]),
                     negative_prompt=[inputs["n_prompt"]] * int(inputs["batch_size"]),
                     height=inputs["height"],
@@ -110,26 +113,85 @@ class StableDiffusion:
                     max_embeddings_multiples=inputs["max_embeddings_multiples"],
                 ).images
 
-        image_output = []
-        for image in images:
-            with io.BytesIO() as buf:
-                image.save(buf, format="PNG")
-                image_output.append(buf.getvalue())
-
         if inputs["upscaler"] != "":
-            upscaled_images = self.upscaler(
+            uplcaled_images = self.upscale(
-                prompt=inputs["prompt"],
+                base_images=base_images,
-                image=images,
+                model_name="RealESRGAN_x4plus",
-                num_inference_steps=inputs["steps"],
+                scale_factor=4,
-                guidance_scale=0,
+                half_precision=False,
-            ).images
+                tile=700,
-            for image in upscaled_images:
+            )
+            base_images.extend(uplcaled_images)
+
+        image_output = []
+        for image in base_images:
             with io.BytesIO() as buf:
                 image.save(buf, format="PNG")
                 image_output.append(buf.getvalue())
 
         return image_output
 
+    @method()
+    def upscale(
+        self,
+        base_images: list[Image.Image],
+        model_name: str = "RealESRGAN_x4plus",
+        scale_factor: float = 4,
+        half_precision: bool = False,
+        tile: int = 0,
+        tile_pad: int = 10,
+        pre_pad: int = 0,
+    ) -> list[Image.Image]:
+        """
+        Upscales the given images using the given model.
+        https://github.com/xinntao/Real-ESRGAN
+        """
+        import numpy
+        import torch
+        from basicsr.archs.rrdbnet_arch import RRDBNet
+        from PIL import Image
+        from realesrgan import RealESRGANer
+        from tqdm import tqdm
+
+        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,
+        )
+
+        torch.cuda.empty_cache()
+        upscaled_imgs = []
+        with tqdm(total=len(base_images)) as progress_bar:
+            for i, img in enumerate(base_images):
+                img = numpy.array(img)
+                enhance_result = upsampler.enhance(img)[0]
+                upscaled_imgs.append(Image.fromarray(enhance_result))
+                progress_bar.update(1)
+        torch.cuda.empty_cache()
+
+        return upscaled_imgs
+
 
 @stub.local_entrypoint()
 def entrypoint(