Upload folder using huggingface_hub

handler.py

@@ -4,8 +4,8 @@ from pathlib import Path
 from typing import Any, Dict, List
 
 from inference import model_fn, predict_fn
-from internals.util.config import set_hf_cache_dir
-from internals.util.model_downloader import BaseModelDownloader
+from internals.util.config import set_hf_cache_dir, set_model_config
+from internals.util.model_loader import load_model_from_config
 
 
 class EndpointHandler:
@@ -13,27 +13,6 @@ class EndpointHandler:
         set_hf_cache_dir(Path.home() / ".cache" / "hf_cache")
         self.model_dir = path
 
-        if os.path.exists(path + "/inference.json"):
-            with open(path + "/inference.json", "r") as f:
-                config = json.loads(f.read())
-                if config.get("model_type") == "huggingface":
-                    self.model_dir = config["model_path"]
-                if config.get("model_type") == "s3":
-                    s3_config = config["model_path"]["s3"]
-                    base_url = s3_config["base_url"]
-
-                    urls = [base_url + item for item in s3_config["paths"]]
-                    out_dir = Path.home() / ".cache" / "base_model"
-                    if out_dir.exists():
-                        print("Model already exist")
-                    else:
-                        print("Downloading model")
-                        BaseModelDownloader(
-                            urls, s3_config["paths"], out_dir
-                        ).download()
-
-                    self.model_dir = str(out_dir)
-
         return model_fn(self.model_dir)
 
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:

inference.py

@@ -21,10 +21,11 @@ from internals.util.avatar import Avatar
 from internals.util.cache import auto_clear_cuda_and_gc, clear_cuda
 from internals.util.commons import download_image, upload_image, upload_images
 from internals.util.config import (get_model_dir, num_return_sequences,
-                                   set_configs_from_task, set_model_dir,
+                                   set_configs_from_task, set_model_config,
                                    set_root_dir)
 from internals.util.failure_hander import FailureHandler
 from internals.util.lora_style import LoraStyle
+from internals.util.model_loader import load_model_from_config
 from internals.util.slack import Slack
 
 torch.backends.cudnn.benchmark = True
@@ -496,13 +497,14 @@ def load_model_by_task(task: Task):
     ):
         text2img_pipe.load(get_model_dir())
         img2img_pipe.create(text2img_pipe)
-        inpainter.create(text2img_pipe)
+        inpainter.load()
         high_res.load(img2img_pipe)
 
         safety_checker.apply(text2img_pipe)
         safety_checker.apply(img2img_pipe)
+        safety_checker.apply(inpainter)
     elif task.get_type() == TaskType.REPLACE_BG:
-        replace_background.load(controlnet=controlnet, high_res=high_res)
+        replace_background.load(inpainter=inpainter, high_res=high_res)
     else:
         if task.get_type() == TaskType.TILE_UPSCALE:
             controlnet.load_tile_upscaler()
@@ -521,7 +523,8 @@ def load_model_by_task(task: Task):
 def model_fn(model_dir):
     print("Logs: model loaded .... starts")
 
-    set_model_dir(model_dir)
+    config = load_model_from_config(model_dir)
+    set_model_config(config)
     set_root_dir(__file__)
 
     FailureHandler.register()

inference2.py

@@ -23,9 +23,10 @@ from internals.util.cache import auto_clear_cuda_and_gc, clear_cuda
 from internals.util.commons import (construct_default_s3_url, upload_image,
                                     upload_images)
 from internals.util.config import (num_return_sequences, set_configs_from_task,
-                                   set_model_dir, set_root_dir)
+                                   set_model_config, set_root_dir)
 from internals.util.failure_hander import FailureHandler
 from internals.util.lora_style import LoraStyle
+from internals.util.model_loader import load_model_from_config
 from internals.util.slack import Slack
 
 torch.backends.cudnn.benchmark = True
@@ -214,7 +215,8 @@ def upscale_image(task: Task):
 def model_fn(model_dir):
     print("Logs: model loaded .... starts")
 
-    set_model_dir(model_dir)
+    config = load_model_from_config(model_dir)
+    set_model_config(config)
     set_root_dir(__file__)
 
     FailureHandler.register()

internals/pipelines/controlnets.py

@@ -4,15 +4,11 @@ import cv2
 import numpy as np
 import torch
 from controlnet_aux import HEDdetector, LineartDetector, OpenposeDetector
-from diffusers import (
-    ControlNetModel,
-    DiffusionPipeline,
-    StableDiffusionControlNetPipeline,
-    UniPCMultistepScheduler,
-)
-from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import (
-    MultiControlNetModel,
-)
+from diffusers import (ControlNetModel, DiffusionPipeline,
+                       StableDiffusionControlNetPipeline,
+                       UniPCMultistepScheduler)
+from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import \
+    MultiControlNetModel
 from PIL import Image
 from torch.nn import Linear
 from tqdm import gui
@@ -22,9 +18,8 @@ import internals.util.image as ImageUtil
 from external.midas import apply_midas
 from internals.data.result import Result
 from internals.pipelines.commons import AbstractPipeline
-from internals.pipelines.tileUpscalePipeline import (
-    StableDiffusionControlNetImg2ImgPipeline,
-)
+from internals.pipelines.tileUpscalePipeline import \
+    StableDiffusionControlNetImg2ImgPipeline
 from internals.util.cache import clear_cuda_and_gc
 from internals.util.commons import download_image
 from internals.util.config import get_hf_cache_dir, get_hf_token, get_model_dir
@@ -86,7 +81,7 @@ class ControlNet(AbstractPipeline):
         if self.__current_task_name == "pose":
             return
         pose = ControlNetModel.from_pretrained(
-            "lllyasviel/sd-controlnet-openpose",
+            "lllyasviel/control_v11p_sd15_openpose",
             torch_dtype=torch.float16,
             cache_dir=get_hf_cache_dir(),
         ).to("cuda")

internals/pipelines/inpainter.py

@@ -5,18 +5,28 @@ from diffusers import StableDiffusionInpaintPipeline
 
 from internals.pipelines.commons import AbstractPipeline
 from internals.util.commons import disable_safety_checker, download_image
-from internals.util.config import get_hf_cache_dir
+from internals.util.config import (get_hf_cache_dir, get_hf_token,
+                                   get_inpaint_model_path)
 
 
 class InPainter(AbstractPipeline):
+    __loaded = False
+
     def load(self):
+        if self.__loaded:
+            return
+
         self.pipe = StableDiffusionInpaintPipeline.from_pretrained(
-            "jayparmr/icbinp_v8_inpaint_v2",
+            get_inpaint_model_path(),
             torch_dtype=torch.float16,
             cache_dir=get_hf_cache_dir(),
+            use_auth_token=get_hf_token(),
         ).to("cuda")
+
         disable_safety_checker(self.pipe)
 
+        self.__loaded = True
+
     def create(self, pipeline: AbstractPipeline):
         self.pipe = StableDiffusionInpaintPipeline(**pipeline.pipe.components).to(
             "cuda"

internals/pipelines/replace_background.py

@@ -2,6 +2,7 @@ from io import BytesIO
 from typing import List, Optional, Union
 
 import torch
+from cv2 import inpaint
 from diffusers import (ControlNetModel,
                        StableDiffusionControlNetInpaintPipeline,
                        StableDiffusionInpaintPipeline, UniPCMultistepScheduler)
@@ -12,10 +13,12 @@ from internals.data.result import Result
 from internals.pipelines.commons import AbstractPipeline
 from internals.pipelines.controlnets import ControlNet
 from internals.pipelines.high_res import HighRes
+from internals.pipelines.inpainter import InPainter
 from internals.pipelines.remove_background import RemoveBackgroundV2
 from internals.pipelines.upscaler import Upscaler
 from internals.util.commons import download_image
-from internals.util.config import get_hf_cache_dir, get_model_dir
+from internals.util.config import (get_hf_cache_dir, get_hf_token,
+                                   get_inpaint_model_path, get_model_dir)
 
 
 class ReplaceBackground(AbstractPipeline):
@@ -25,7 +28,7 @@ class ReplaceBackground(AbstractPipeline):
         self,
         upscaler: Optional[Upscaler] = None,
         remove_background: Optional[RemoveBackgroundV2] = None,
-        controlnet: Optional[ControlNet] = None,
+        inpainter: Optional[InPainter] = None,
         high_res: Optional[HighRes] = None,
     ):
         if self.__loaded:
@@ -35,18 +38,19 @@ class ReplaceBackground(AbstractPipeline):
             torch_dtype=torch.float16,
             cache_dir=get_hf_cache_dir(),
         ).to("cuda")
-        if controlnet:
-            controlnet.load_linearart()
+        if inpainter:
+            inpainter.load()
             pipe = StableDiffusionControlNetInpaintPipeline(
-                **controlnet.pipe.components
+                **inpainter.pipe.components,
+                controlnet=controlnet_model,
             )
-            pipe.controlnet = controlnet_model
         else:
             pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
                 "runwayml/stable-diffusion-inpainting",
                 controlnet=controlnet_model,
                 torch_dtype=torch.float16,
                 cache_dir=get_hf_cache_dir(),
+                use_auth_token=get_hf_token(),
             )
         pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
         pipe.to("cuda")
@@ -104,14 +108,14 @@ class ReplaceBackground(AbstractPipeline):
 
         print(width, height, n_width, n_height)
 
+        resolution = min(n_width, n_height)
         if extend_object:
-            condition_image = ControlNet.linearart_condition_image(image).resize(
-                (n_width, n_height)
-            )
+            condition_image = ControlNet.linearart_condition_image(image)
+            condition_image = ImageUtil.resize_image(condition_image, resolution)
             condition_image = ImageUtil.padd_image(condition_image, width, height)
             condition_image = condition_image.convert("RGB")
 
-        image = image.resize((n_width, n_height))
+        image = ImageUtil.resize_image(image, resolution)
         image = ImageUtil.padd_image(image, width, height)
 
         mask = image.copy()
@@ -130,46 +134,20 @@ class ReplaceBackground(AbstractPipeline):
             condition_image = ControlNet.linearart_condition_image(image)
         mask = mask.convert("RGB")
 
-        if apply_high_res and hasattr(self, "high_res"):
-            w, h = HighRes.get_intermediate_dimension(width, height)
-            result = self.pipe.__call__(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                image=image,
-                mask_image=mask,
-                control_image=condition_image,
-                controlnet_conditioning_scale=conditioning_scale,
-                guidance_scale=9,
-                strength=1,
-                num_inference_steps=steps,
-                height=w,
-                width=h,
-            )
-            for i, _ in enumerate(result.images):
-                out_bytes = self.upscaler.upscale(
-                    image=result.images[i],
-                    width=w,
-                    height=h,
-                    face_enhance=False,
-                    resize_dimension=max(width, height),
-                )
-                result.images[i] = Image.open(BytesIO(out_bytes)).convert("RGB")
-            result = Result.from_result(result)
-        else:
-            result = self.pipe.__call__(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                image=image,
-                mask_image=mask,
-                control_image=condition_image,
-                controlnet_conditioning_scale=conditioning_scale,
-                guidance_scale=9,
-                strength=1,
-                height=height,
-                num_inference_steps=steps,
-                width=width,
-            )
-            result = Result.from_result(result)
+        result = self.pipe.__call__(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=image,
+            mask_image=mask,
+            control_image=condition_image,
+            controlnet_conditioning_scale=conditioning_scale,
+            guidance_scale=9,
+            strength=1,
+            height=height,
+            num_inference_steps=steps,
+            width=width,
+        )
+        result = Result.from_result(result)
 
         images, has_nsfw = result
 

internals/util/config.py

@@ -3,13 +3,14 @@ from pathlib import Path
 from typing import Union
 
 from internals.data.task import Task
+from internals.util.model_loader import ModelConfig
 
 env = "prod"
 nsfw_threshold = 0.0
 nsfw_access = False
 access_token = ""
 root_dir = ""
-model_dir = ""
+model_config = None
 hf_token = "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"
 hf_cache_dir = "/tmp/hf_hub"
 
@@ -28,16 +29,16 @@ def get_hf_cache_dir():
     return hf_cache_dir
 
 
-def set_model_dir(dir: str):
-    global model_dir
-    model_dir = dir
-
-
 def set_root_dir(main_file: str):
     global root_dir
     root_dir = os.path.dirname(os.path.abspath(main_file))
 
 
+def set_model_config(config: ModelConfig):
+    global model_config
+    model_config = config
+
+
 def set_configs_from_task(task: Task):
     global env, nsfw_threshold, nsfw_access, access_token
     name = task.get_queue_name()
@@ -51,8 +52,13 @@ def set_configs_from_task(task: Task):
 
 
 def get_model_dir():
-    global model_dir
-    return model_dir
+    global model_config
+    return model_config.base_model_path  # pyright: ignore
+
+
+def get_inpaint_model_path():
+    global model_config
+    return model_config.base_inpaint_model_path  # pyright: ignore
 
 
 def get_root_dir():

internals/util/model_loader.py

@@ -0,0 +1,187 @@
+import json
+import os
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+from threading import Thread
+from typing import Any, Dict, List, Optional
+
+import requests
+from tqdm import tqdm
+
+
+@dataclass
+class ModelConfig:
+    base_model_path: str
+    base_inpaint_model_path: str
+
+
+def load_model_from_config(path):
+    m_config = ModelConfig(path, path)
+    if os.path.exists(path + "/inference.json"):
+        with open(path + "/inference.json", "r") as f:
+            config = json.loads(f.read())
+            model_path = config.get("model_path", path)
+            inpaint_model_path = config.get("inpaint_model_path", path)
+
+            m_config.base_model_path = model_path
+            m_config.base_inpaint_model_path = inpaint_model_path
+
+            #
+            # if config.get("model_type") == "huggingface":
+            #     model_dir = config["model_path"]
+            # if config.get("model_type") == "s3":
+            #     s3_config = config["model_path"]["s3"]
+            #     base_url = s3_config["base_url"]
+            #
+            #     urls = [base_url + item for item in s3_config["paths"]]
+            #     out_dir = Path.home() / ".cache" / "base_model"
+            #     if out_dir.exists():
+            #         print("Model already exist")
+            #     else:
+            #         print("Downloading model")
+            #         BaseModelDownloader(urls, s3_config["paths"], out_dir).download()
+            #     model_dir = str(out_dir)
+    return m_config
+
+
+class BaseModelDownloader:
+    """
+    A utility for fast download of base model from S3 or any CDN served storage.
+    Works by downloading multiple files in parallel and dividing large files
+    into smaller chunks and combining them at the end.
+
+    Currently it uses multithreading (not multiprocessing) assuming GIL won't
+    interfere with network/disk IO.
+
+    Created by: KP
+    """
+
+    def __init__(self, urls: List[str], url_paths: List[str], out_dir: Path):
+        self.urls = urls
+        self.url_paths = url_paths
+        shutil.rmtree(out_dir, ignore_errors=True)
+        out_dir.mkdir(parents=True, exist_ok=True)
+        self.out_dir = out_dir
+
+    def download(self):
+        threads = []
+        batch_urls = {}
+
+        for url, url_path in zip(self.urls, self.url_paths):
+            out_dir = self.out_dir / url_path
+            self.out_dir.parent.mkdir(parents=True, exist_ok=True)
+            if url.endswith(".bin"):
+                if "unet/" in url_path:
+                    thread = Thread(
+                        target=self.__download_parallel, args=(url, out_dir, 6)
+                    )
+                    thread.start()
+                    threads.append(thread)
+                else:
+                    thread = Thread(
+                        target=self.__download_files, args=([url], [out_dir])
+                    )
+                    thread.start()
+                    threads.append(thread)
+                    pass
+            else:
+                batch_urls[url] = out_dir
+
+        if batch_urls:
+            thread = Thread(
+                target=self.__download_files,
+                args=(list(batch_urls.keys()), list(batch_urls.values())),
+            )
+            thread.start()
+            threads.append(thread)
+            pass
+
+        for thread in threads:
+            thread.join()
+
+    def __download_parallel(self, url, output_filename, num_parts=4):
+        response = requests.head(url)
+        total_size = int(response.headers.get("content-length", 0))
+        print("total_size", total_size)
+
+        chunk_size = total_size // num_parts
+        ranges = [
+            (i * chunk_size, (i + 1) * chunk_size - 1) for i in range(num_parts - 1)
+        ]
+        ranges.append((ranges[-1][1] + 1, total_size))
+
+        print(ranges)
+
+        save_dir = Path.home() / ".cache" / "download_parts"
+        os.makedirs(save_dir, exist_ok=True)
+
+        threads = []
+        for i, (start, end) in enumerate(ranges):
+            thread = Thread(
+                target=self.__download_part, args=(url, start, end, i, save_dir)
+            )
+            thread.start()
+            threads.append(thread)
+
+        for thread in threads:
+            thread.join()
+
+        self.__combine_parts(save_dir, output_filename, num_parts)
+        os.rmdir(save_dir)
+
+    def __combine_parts(self, save_dir, output_filename, num_parts):
+        part_files = [os.path.join(save_dir, f"part_{i}.tmp") for i in range(num_parts)]
+
+        output_filename.parent.mkdir(parents=True, exist_ok=True)
+        with open(output_filename, "wb") as output_file:
+            for part_file in part_files:
+                print("combining: ", part_file)
+                with open(part_file, "rb") as part:
+                    output_file.write(part.read())
+
+            out_file_size = output_file.tell()
+            print("out_file_size", out_file_size)
+
+        for part_file in part_files:
+            os.remove(part_file)
+
+    def __download_part(self, url, start_byte, end_byte, part_num, save_dir):
+        headers = {"Range": f"bytes={start_byte}-{end_byte}"}
+        response = requests.get(url, headers=headers, stream=True)
+
+        part_filename = os.path.join(save_dir, f"part_{part_num}.tmp")
+        print("Downloading part: ", url, part_filename, end_byte - start_byte)
+
+        with open(part_filename, "wb") as part_file, tqdm(
+            desc=str(part_filename),
+            total=end_byte - start_byte,
+            unit="B",
+            unit_scale=True,
+            unit_divisor=1024,
+        ) as bar:
+            for chunk in response.iter_content(chunk_size=8192):
+                if chunk:
+                    size = part_file.write(chunk)
+                    bar.update(size)
+
+        return part_filename
+
+    def __download_files(self, urls, out_paths: List[Path]):
+        for url, out_path in zip(urls, out_paths):
+            out_path.parent.mkdir(parents=True, exist_ok=True)
+            with requests.get(url, stream=True) as r:
+                print("Downloading: ", url)
+                total_size = int(r.headers.get("content-length", 0))
+                chunk_size = 8192
+                r.raise_for_status()
+                with open(out_path, "wb") as f, tqdm(
+                    desc=str(out_path),
+                    total=total_size,
+                    unit="B",
+                    unit_scale=True,
+                    unit_divisor=1024,
+                ) as bar:
+                    for data in r.iter_content(chunk_size=chunk_size):
+                        size = f.write(data)
+                        bar.update(size)

requirements.txt

@@ -5,7 +5,7 @@ fastapi==0.87.0
 Pillow==9.3.0
 redis==4.3.4
 requests==2.28.1
-transformers
+transformers==4.34.1
 rembg==2.0.30
 gfpgan==1.3.8
 rembg==2.0.30