Upload folder using huggingface_hub

handler.py

@@ -1,11 +1,38 @@
+import json
+import os
+from pathlib import Path
 from typing import Any, Dict, List
 
 from inference import model_fn, predict_fn
+from internals.util.model_downloader import BaseModelDownloader
 
 
 class EndpointHandler:
     def __init__(self, path=""):
-        return model_fn(path)
+        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]]]:
         return predict_fn(data, None)

inference.py

@@ -15,18 +15,13 @@ from internals.pipelines.safety_checker import SafetyChecker
 from internals.util.anomaly import remove_colors
 from internals.util.args import apply_style_args
 from internals.util.avatar import Avatar
-from internals.util.cache import auto_clear_cuda_and_gc, clear_cuda, clear_cuda_and_gc
-from internals.util.commons import (
-    download_image,
-    pickPoses,
-    upload_image,
-    upload_images,
-)
-from internals.util.config import (
-    num_return_sequences,
-    set_configs_from_task,
-    set_root_dir,
-)
+from internals.util.cache import (auto_clear_cuda_and_gc, clear_cuda,
+                                  clear_cuda_and_gc)
+from internals.util.commons import (download_image, pickPoses, upload_image,
+                                    upload_images)
+from internals.util.config import (get_model_dir, num_return_sequences,
+                                   set_configs_from_task, set_model_dir,
+                                   set_root_dir)
 from internals.util.failure_hander import FailureHandler
 from internals.util.lora_style import LoraStyle
 from internals.util.slack import Slack
@@ -442,27 +437,48 @@ def inpaint(task: Task):
     return {"modified_prompts": prompt, "generated_image_urls": generated_image_urls}
 
 
+def load_model_by_task(task: Task):
+    if (
+        task.get_type()
+        in [
+            TaskType.TEXT_TO_IMAGE,
+            TaskType.IMAGE_TO_IMAGE,
+            TaskType.INPAINT,
+        ]
+        and not text2img_pipe.is_loaded()
+    ):
+        text2img_pipe.load(get_model_dir())
+        img2img_pipe.create(text2img_pipe)
+        inpainter.create(text2img_pipe)
+
+        safety_checker.apply(text2img_pipe)
+        safety_checker.apply(img2img_pipe)
+    else:
+        if task.get_type() == TaskType.TILE_UPSCALE:
+            controlnet.load_tile_upscaler()
+        elif task.get_type() == TaskType.CANNY:
+            controlnet.load_canny()
+        elif task.get_type() == TaskType.SCRIBBLE:
+            controlnet.load_scribble()
+        elif task.get_type() == TaskType.LINEARART:
+            controlnet.load_linearart()
+        elif task.get_type() == TaskType.POSE:
+            controlnet.load_pose()
+
+        safety_checker.apply(controlnet)
+
+
 def model_fn(model_dir):
     print("Logs: model loaded .... starts")
 
+    set_model_dir(model_dir)
     set_root_dir(__file__)
 
     FailureHandler.register()
 
     avatar.load_local(model_dir)
 
-    prompt_modifier.load()
-    pose_detector.load()
-    img2text.load()
-    img_classifier.load()
-
     lora_style.load(model_dir)
-    safety_checker.load()
-
-    controlnet.load(model_dir)
-    text2img_pipe.load(model_dir)
-    img2img_pipe.create(text2img_pipe)
-    inpainter.create(text2img_pipe)
 
     print("Logs: model loaded ....")
     return
@@ -479,10 +495,8 @@ def predict_fn(data, pipe):
         # Set set_environment
         set_configs_from_task(task)
 
-        # Apply safety checkers based on environment
-        safety_checker.apply(text2img_pipe)
-        safety_checker.apply(img2img_pipe)
-        safety_checker.apply(controlnet)
+        # Load model based on task
+        load_model_by_task(task)
 
         # Apply arguments
         apply_style_args(data)
@@ -497,10 +511,10 @@ def predict_fn(data, pipe):
 
         if task_type == TaskType.TEXT_TO_IMAGE:
             # character sheet
-            if "character sheet" in task.get_prompt().lower():
-                return pose(task, s3_outkey="", poses=pickPoses())
-            else:
-                return text2img(task)
+            # if "character sheet" in task.get_prompt().lower():
+            #     return pose(task, s3_outkey="", poses=pickPoses())
+            # else:
+            return text2img(task)
         elif task_type == TaskType.IMAGE_TO_IMAGE:
             return img2img(task)
         elif task_type == TaskType.CANNY:

inference2.py

@@ -7,17 +7,18 @@ from internals.data.task import ModelType, Task, TaskType
 from internals.pipelines.inpainter import InPainter
 from internals.pipelines.object_remove import ObjectRemoval
 from internals.pipelines.prompt_modifier import PromptModifier
-from internals.pipelines.remove_background import (RemoveBackground,
-                                                   RemoveBackgroundV2)
+from internals.pipelines.remove_background import RemoveBackground, RemoveBackgroundV2
 from internals.pipelines.replace_background import ReplaceBackground
 from internals.pipelines.safety_checker import SafetyChecker
 from internals.pipelines.upscaler import Upscaler
 from internals.util.avatar import Avatar
 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_root_dir)
+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_root_dir,
+)
 from internals.util.failure_hander import FailureHandler
 from internals.util.slack import Slack
 
@@ -189,6 +190,7 @@ def predict_fn(data, pipe):
 
         # Apply safety checker based on environment
         safety_checker.apply(inpainter)
+        safety_checker.apply(replace_background)
 
         # Fetch avatars
         avatar.fetch_from_network(task.get_model_id())

internals/data/task.py

@@ -30,7 +30,7 @@ class Task:
     def __init__(self, data):
         self.__data = data
         if data.get("seed", -1) == None or self.get_seed() == -1:
-            self.__data["seed"] = np.random.randint(0, np.iinfo(np.int64).max)
+            self.__data["seed"] = np.random.randint(0, np.iinfo(np.int32).max)
         prompt = data.get("prompt", "")
         if prompt is None:
             self.__data["prompt"] = ""

internals/pipelines/commons.py

@@ -7,7 +7,7 @@ from diffusers import StableDiffusionImg2ImgPipeline
 from internals.data.result import Result
 from internals.pipelines.twoStepPipeline import two_step_pipeline
 from internals.util.commons import disable_safety_checker, download_image
-from internals.util.config import num_return_sequences
+from internals.util.config import get_hf_token, num_return_sequences
 
 
 class AbstractPipeline:
@@ -28,10 +28,15 @@ class Text2Img(AbstractPipeline):
 
     def load(self, model_dir: str):
         self.pipe = two_step_pipeline.from_pretrained(
-            model_dir, torch_dtype=torch.float16
+            model_dir, torch_dtype=torch.float16, use_auth_token=get_hf_token()
         ).to("cuda")
         self.__patch()
 
+    def is_loaded(self):
+        if hasattr(self, "pipe"):
+            return True
+        return False
+
     def create(self, pipeline: AbstractPipeline):
         self.pipe = two_step_pipeline(**pipeline.pipe.components).to("cuda")
         self.__patch()
@@ -115,7 +120,7 @@ class Text2Img(AbstractPipeline):
 class Img2Img(AbstractPipeline):
     def load(self, model_dir: str):
         self.pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
-            model_dir, torch_dtype=torch.float16
+            model_dir, torch_dtype=torch.float16, use_auth_token=get_hf_token()
         ).to("cuda")
         self.__patch()
 

internals/pipelines/controlnets.py

@@ -4,33 +4,43 @@ import cv2
 import numpy as np
 import torch
 from controlnet_aux import HEDdetector, LineartDetector, OpenposeDetector
-from diffusers import (ControlNetModel, DiffusionPipeline,
-                       StableDiffusionControlNetPipeline,
-                       UniPCMultistepScheduler)
+from diffusers import (
+    ControlNetModel,
+    DiffusionPipeline,
+    StableDiffusionControlNetPipeline,
+    UniPCMultistepScheduler,
+)
 from PIL import Image
 from torch.nn import Linear
 from tqdm import gui
 
 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_token, get_model_dir
 
 
 class ControlNet(AbstractPipeline):
     __current_task_name = ""
+    __loaded = False
 
-    def load(self, model_dir: str):
-        # we will load canny by default
-        self.load_scribble()
+    def load(self):
+        if self.__loaded:
+            return
+
+        if not self.controlnet:
+            self.load_pose()
 
         # controlnet pipeline for tile upscaler
         pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
-            model_dir,
+            get_model_dir(),
             controlnet=self.controlnet,
             torch_dtype=torch.float16,
+            use_auth_token=get_hf_token(),
         ).to("cuda")
         # pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
         pipe.enable_model_cpu_offload()
@@ -43,6 +53,8 @@ class ControlNet(AbstractPipeline):
         pipe2.enable_xformers_memory_efficient_attention()
         self.pipe2 = pipe2
 
+        self.__loaded = True
+
     def load_canny(self):
         if self.__current_task_name == "canny":
             return
@@ -51,6 +63,9 @@ class ControlNet(AbstractPipeline):
         ).to("cuda")
         self.__current_task_name = "canny"
         self.controlnet = canny
+
+        self.load()
+
         if hasattr(self, "pipe"):
             self.pipe.controlnet = canny
         if hasattr(self, "pipe2"):
@@ -65,6 +80,9 @@ class ControlNet(AbstractPipeline):
         ).to("cuda")
         self.__current_task_name = "pose"
         self.controlnet = pose
+
+        self.load()
+
         if hasattr(self, "pipe"):
             self.pipe.controlnet = pose
         if hasattr(self, "pipe2"):
@@ -79,6 +97,9 @@ class ControlNet(AbstractPipeline):
         ).to("cuda")
         self.__current_task_name = "tile_upscaler"
         self.controlnet = tile_upscaler
+
+        self.load()
+
         if hasattr(self, "pipe"):
             self.pipe.controlnet = tile_upscaler
         if hasattr(self, "pipe2"):
@@ -93,6 +114,9 @@ class ControlNet(AbstractPipeline):
         ).to("cuda")
         self.__current_task_name = "scribble"
         self.controlnet = scribble
+
+        self.load()
+
         if hasattr(self, "pipe"):
             self.pipe.controlnet = scribble
         if hasattr(self, "pipe2"):
@@ -108,6 +132,9 @@ class ControlNet(AbstractPipeline):
         ).to("cuda")
         self.__current_task_name = "linearart"
         self.controlnet = linearart
+
+        self.load()
+
         if hasattr(self, "pipe"):
             self.pipe.controlnet = linearart
         if hasattr(self, "pipe2"):

internals/pipelines/img_classifier.py

@@ -6,16 +6,24 @@ from internals.util.commons import download_image
 
 
 class ImageClassifier:
+    __loaded = False
+
     def __init__(self, candidates: List[str] = ["realistic", "anime", "comic"]):
         self.__candidates = candidates
 
     def load(self):
+        if self.__loaded:
+            return
+
         self.pipe = pipeline(
             "zero-shot-image-classification",
             model="philschmid/clip-zero-shot-image-classification",
         )
 
+        self.__loaded = True
+
     def classify(self, image_url: str, width: int, height: int) -> str:
+        self.load()
         image = download_image(image_url).resize((width, height))
         results = self.pipe.__call__([image], candidate_labels=self.__candidates)
         results = results[0]

internals/pipelines/img_to_text.py

@@ -8,7 +8,12 @@ from internals.util.commons import download_image
 
 
 class Image2Text:
+    __loaded = False
+
     def load(self):
+        if self.__loaded:
+            return
+
         self.processor = BlipProcessor.from_pretrained(
             "Salesforce/blip-image-captioning-large"
         )
@@ -16,7 +21,10 @@ class Image2Text:
             "Salesforce/blip-image-captioning-large", torch_dtype=torch.float16
         ).to("cuda")
 
+        self.__loaded = True
+
     def process(self, imageUrl: str) -> str:
+        self.load()
         image = download_image(imageUrl).resize((512, 512))
         inputs = self.processor.__call__(image, return_tensors="pt").to(
             "cuda", torch.float16

internals/pipelines/pose_detector.py

@@ -36,6 +36,7 @@ class PoseDetector:
         client_coordinates: Optional[dict],
     ) -> Image.Image:
         "Infer pose coordinates from image, map head and body coordinates to infered ones, create pose"
+        self.load()
         if type(image) is str:
             image = download_image(image)
 
@@ -103,6 +104,7 @@ class PoseDetector:
         return image
 
     def infer(self, image: Union[str, Image.Image], width, height) -> dict:
+        self.load()
         candidate = []
         subset = []
 

internals/pipelines/prompt_modifier.py

@@ -4,11 +4,15 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
 
 
 class PromptModifier:
+    __loaded = False
+
     def __init__(self, num_of_sequences: Optional[int] = 4):
         self.__blacklist = {"alphonse mucha": "", "adolphe bouguereau": ""}
         self.__num_of_sequences = num_of_sequences
 
     def load(self):
+        if self.__loaded:
+            return
         self.prompter_model = AutoModelForCausalLM.from_pretrained(
             "Gustavosta/MagicPrompt-Stable-Diffusion"
         )
@@ -18,7 +22,10 @@ class PromptModifier:
         self.prompter_tokenizer.pad_token = self.prompter_tokenizer.eos_token
         self.prompter_tokenizer.padding_side = "left"
 
+        self.__loaded = True
+
     def modify(self, text: str, num_of_sequences: Optional[int] = None) -> List[str]:
+        self.load()
         eos_id = self.prompter_tokenizer.eos_token_id
         # restricted_words_list = ["octane", "cyber"]
         # restricted_words_token_ids = prompter_tokenizer(

internals/pipelines/replace_background.py

@@ -12,13 +12,14 @@ from PIL import Image, ImageFilter, ImageOps
 
 import internals.util.image as ImageUtil
 from internals.data.result import Result
+from internals.pipelines.commons import AbstractPipeline
 from internals.pipelines.controlnets import ControlNet
 from internals.pipelines.remove_background import RemoveBackgroundV2
 from internals.pipelines.upscaler import Upscaler
 from internals.util.commons import download_image
 
 
-class ReplaceBackground:
+class ReplaceBackground(AbstractPipeline):
     def load(self, upscaler: Upscaler, remove_background: RemoveBackgroundV2):
         controlnet = ControlNetModel.from_pretrained(
             "lllyasviel/control_v11p_sd15_lineart", torch_dtype=torch.float16

internals/pipelines/safety_checker.py

@@ -18,13 +18,24 @@ def cosine_distance(image_embeds, text_embeds):
 
 
 class SafetyChecker:
+    __loaded = False
+
     def load(self):
+        if self.__loaded:
+            return
+
         self.model = StableDiffusionSafetyCheckerV2.from_pretrained(
             "CompVis/stable-diffusion-safety-checker", torch_dtype=torch.float16
         ).to("cuda")
 
+        self.__loaded = True
+
     def apply(self, pipeline: AbstractPipeline):
+        self.load()
+
         model = self.model if not get_nsfw_access() else None
+        if not pipeline:
+            return
         if hasattr(pipeline, "pipe"):
             pipeline.pipe.safety_checker = model
         if hasattr(pipeline, "pipe2"):

internals/util/config.py

@@ -7,10 +7,17 @@ nsfw_threshold = 0.0
 nsfw_access = False
 access_token = ""
 root_dir = ""
+model_dir = ""
+hf_token = "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"
 
 num_return_sequences = 4  # the number of results to generate
 
 
+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))
@@ -28,6 +35,11 @@ def set_configs_from_task(task: Task):
     access_token = task.get_access_token()
 
 
+def get_model_dir():
+    global model_dir
+    return model_dir
+
+
 def get_root_dir():
     global root_dir
     return root_dir
@@ -48,6 +60,11 @@ def get_nsfw_access():
     return nsfw_access
 
 
+def get_hf_token():
+    global hf_token
+    return hf_token
+
+
 def api_headers():
     return {
         "Access-Token": access_token,

internals/util/model_downloader.py

@@ -0,0 +1,151 @@
+import json
+import os
+import shutil
+from pathlib import Path
+from threading import Thread
+from typing import Any, Dict, List
+
+import requests
+from tqdm import tqdm
+
+
+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

@@ -22,20 +22,19 @@ kornia==0.5.0
 pytorch-lightning==1.2.9
 mmpose==0.29.0
 mmdet==2.28.2
-mmengine
-pydash
-scikit-learn
-accelerate
-pandas
-xformers
-torchvision
-scikit-image
-omegaconf
-webdataset
-git+https://github.com/cloneofsimo/lora.git
+https://comic-assets.s3.ap-south-1.amazonaws.com/packages/v1/lora-diffusion-0.1.7.zip
+mmengine==0.8.4
+pydash==7.0.6
+scikit-learn==1.3.0
+accelerate==0.22.0
+pandas==2.0.3
+xformers==0.0.21
+scikit-image==0.19.3
+omegaconf==2.3.0
+webdataset==0.2.48
 https://comic-assets.s3.ap-south-1.amazonaws.com/packages/mmcv_full-1.7.0-cp39-cp39-linux_x86_64.whl
 python-dateutil==2.8.2
 PyYAML
-torchvision
-imgaug
-tqdm
+torchvision==0.15.2
+imgaug==0.4.0
+tqdm==4.64.1