Upload folder using huggingface_hub

inference.py

@@ -727,9 +727,6 @@ def load_model_by_task(task_type: TaskType, model_id=-1):
         inpainter.init(text2img_pipe)
         controlnet.init(text2img_pipe)
 
-        safety_checker.apply(text2img_pipe)
-        safety_checker.apply(img2img_pipe)
-
     if task_type == TaskType.INPAINT:
         inpainter.load()
         safety_checker.apply(inpainter)
@@ -756,7 +753,11 @@ def load_model_by_task(task_type: TaskType, model_id=-1):
         elif task_type == TaskType.POSE:
             controlnet.load_model("pose")
 
-        safety_checker.apply(controlnet)
+
+def apply_safety_checkers():
+    safety_checker.apply(text2img_pipe)
+    safety_checker.apply(img2img_pipe)
+    safety_checker.apply(controlnet)
 
 
 def model_fn(model_dir):
@@ -798,6 +799,9 @@ def predict_fn(data, pipe):
             task.get_type() or TaskType.TEXT_TO_IMAGE, task.get_model_id()
         )
 
+        # Apply safety checkers
+        apply_safety_checkers()
+
         # Realtime generation apis
         if task_type == TaskType.RT_DRAW_SEG:
             return rt_draw_seg(task)
@@ -814,10 +818,6 @@ def predict_fn(data, pipe):
         avatar.fetch_from_network(task.get_model_id())
 
         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)
         elif task_type == TaskType.IMAGE_TO_IMAGE:
             return img2img(task)

internals/pipelines/realtime_draw.py

@@ -72,10 +72,16 @@ class RealtimeDraw(AbstractPipeline):
         torch.manual_seed(seed)
 
         if not image:
-            image = Image.new("RGB", (512, 512), color=0)
+            size = (512, 512)
+            if image2:
+                size = image2.size
+            image = Image.new("RGB", size, color=0)
 
         if not image2:
-            image2 = Image.new("RGB", image.size, color=0)
+            size = (512, 512)
+            if image:
+                size = image.size
+            image2 = Image.new("RGB", size, color=0)
 
         image = ImageUtil.resize_image(image, 512)
 
@@ -91,6 +97,8 @@ class RealtimeDraw(AbstractPipeline):
             negative_prompt=negative_prompt,
             guidance_scale=10,
             strength=0.9,
+            width=image.size[0],
+            height=image.size[1],
             controlnet_conditioning_scale=[1.0, 0.8],
         ).images[0]