Upload folder using huggingface_hub

inference2.py

@@ -4,7 +4,7 @@ from io import BytesIO
 import torch
 
 import internals.util.prompt as prompt_util
-from internals.data.dataAccessor import update_db
+from internals.data.dataAccessor import update_db, update_db_source_failed
 from internals.data.task import ModelType, Task, TaskType
 from internals.pipelines.controlnets import ControlNet
 from internals.pipelines.high_res import HighRes
@@ -194,7 +194,10 @@ def replace_bg(task: Task):
 def upscale_image(task: Task):
     output_key = "crecoAI/{}_upscale.png".format(task.get_taskId())
     out_img = None
-    if task.get_modelType() == ModelType.ANIME:
+    if (
+        task.get_modelType() == ModelType.ANIME
+        or task.get_modelType() == ModelType.COMIC
+    ):
         print("Using Anime model")
         out_img = upscaler.upscale_anime(
             image=task.get_imageUrl(),
@@ -297,4 +300,5 @@ def predict_fn(data, pipe):
         print(f"Error: {e}")
         slack.error_alert(task, e)
         controlnet.cleanup()
+        update_db_source_failed(task.get_sourceId(), task.get_userId())
         return None

internals/data/dataAccessor.py

@@ -19,7 +19,6 @@ class RetryRequest:
 
 
 def updateSource(sourceId, userId, state):
-    print("update source is called")
     url = api_endpoint() + f"/autodraft-crecoai/source/{sourceId}"
     headers = {
         "Content-Type": "application/json",
@@ -32,7 +31,6 @@ def updateSource(sourceId, userId, state):
     try:
         with RetryRequest() as session:
             response = session.patch(url, headers=headers, json=data, timeout=10)
-        print("update source response", response)
     except requests.exceptions.Timeout:
         print("Request timed out while updating source")
     except requests.exceptions.RequestException as e:
@@ -42,7 +40,6 @@ def updateSource(sourceId, userId, state):
 
 
 def saveGeneratedImages(sourceId, userId, has_nsfw: bool):
-    print("save generation called")
     url = (
         api_endpoint()
         + "/autodraft-crecoai/source/"
@@ -70,7 +67,6 @@ def saveGeneratedImages(sourceId, userId, has_nsfw: bool):
 
 def getStyles() -> Optional[Dict]:
     url = api_endpoint() + "/autodraft-crecoai/style"
-    print(url)
     try:
         with RetryRequest() as session:
             response = session.get(

internals/data/task.py

@@ -26,6 +26,10 @@ class ModelType(Enum):
     ANIME = 10001
     COMIC = 10002
 
+    @classmethod
+    def _missing_(cls, value):
+        return cls.REAL
+
 
 class Task:
     def __init__(self, data):
@@ -156,6 +160,9 @@ class Task:
     def get_high_res_fix(self) -> bool:
         return self.__data.get("high_res_fix", False)
 
+    def get_base_dimension(self) -> int:
+        return self.__data.get("base_dimension", 512)
+
     def get_raw(self) -> dict:
         return self.__data.copy()
 

internals/pipelines/high_res.py

@@ -5,7 +5,7 @@ from PIL import Image
 
 from internals.data.result import Result
 from internals.pipelines.commons import AbstractPipeline, Img2Img
-from internals.util.config import get_model_dir
+from internals.util.config import get_model_dir, get_base_dimension
 
 
 class HighRes(AbstractPipeline):
@@ -42,7 +42,7 @@ class HighRes(AbstractPipeline):
 
     @staticmethod
     def get_intermediate_dimension(target_width: int, target_height: int):
-        def_size = 1024
+        def_size = get_base_dimension()
 
         desired_pixel_count = def_size * def_size
         actual_pixel_count = target_width * target_height

internals/pipelines/upscaler.py

@@ -87,7 +87,7 @@ class Upscaler:
             num_in_ch=3,
             num_out_ch=3,
             num_feat=64,
-            num_block=23,
+            num_block=6,
             num_grow_ch=32,
             scale=4,
         )

internals/util/config.py

@@ -13,6 +13,7 @@ root_dir = ""
 model_config = None
 hf_token = "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"
 hf_cache_dir = "/tmp/hf_hub"
+base_dimension = 512  # needed for high res
 
 num_return_sequences = 4  # the number of results to generate
 
@@ -40,7 +41,7 @@ def set_model_config(config: ModelConfig):
 
 
 def set_configs_from_task(task: Task):
-    global env, nsfw_threshold, nsfw_access, access_token
+    global env, nsfw_threshold, nsfw_access, access_token, base_dimension
     name = task.get_queue_name()
     if name.startswith("gamma"):
         env = "gamma"
@@ -49,6 +50,7 @@ def set_configs_from_task(task: Task):
     nsfw_threshold = task.get_nsfw_threshold()
     nsfw_access = task.can_access_nsfw()
     access_token = task.get_access_token()
+    base_dimension = task.get_base_dimension()
 
 
 def get_model_dir():
@@ -61,6 +63,11 @@ def get_inpaint_model_path():
     return model_config.base_inpaint_model_path  # pyright: ignore
 
 
+def get_base_dimension():
+    global base_dimension
+    return base_dimension
+
+
 def get_is_sdxl():
     global model_config
     return model_config.is_sdxl  # pyright: ignore