Upload mod.py

mod.py

@@ -46,6 +46,7 @@ control_scales = [0] * num_cns
 last_model = models[0]
 last_cn_on = False
 
+
 def is_repo_name(s):
     import re
     return re.fullmatch(r'^[^/,\s\"\']+/[^/,\s\"\']+$', s)
@@ -116,6 +117,20 @@ def change_base_model(repo_id: str, cn_on: bool, progress=gr.Progress(track_tqdm
     return gr.update(visible=True)
 
 
+def expand2square(pil_img: Image.Image, background_color: tuple=(0, 0, 0)):
+    width, height = pil_img.size
+    if width == height:
+        return pil_img
+    elif width > height:
+        result = Image.new(pil_img.mode, (width, width), background_color)
+        result.paste(pil_img, (0, (width - height) // 2))
+        return result
+    else:
+        result = Image.new(pil_img.mode, (height, height), background_color)
+        result.paste(pil_img, ((height - width) // 2, 0))
+        return result
+
+
 # https://huggingface.co/spaces/DamarJati/FLUX.1-DEV-Canny/blob/main/app.py
 def resize_image(image, target_width, target_height, crop=True):
     from image_datasets.canny_dataset import c_crop
@@ -169,8 +184,9 @@ def get_control_params():
 
 from preprocessor import Preprocessor
 def preprocess_image(image: Image.Image, control_mode: str, height: int, width: int, preprocess_resolution: int):
+    if control_mode == "None": return image
     image_resolution = max(width, height)
-    image_before = resize_image(image, image_resolution, image_resolution, True)
+    image_before = resize_image(expand2square(image), image_resolution, image_resolution, False)
     # generated control_
     print("start to generate control image")
     preprocessor = Preprocessor()
@@ -213,15 +229,15 @@ def preprocess_image(image: Image.Image, control_mode: str, height: int, width:
             detect_resolution=preprocess_resolution,
         )
     
-    if control_mode == "low_quality" or  control_mode == "gray" or  control_mode == "blur" or control_mode == "tile":
+    if control_mode == "low_quality" or control_mode == "gray" or control_mode == "blur" or control_mode == "tile":
         control_image = image_before
         image_width = 768
         image_height = 768
     else:
         # make sure control image size is same as resized_image
-        image_width, image_height =  control_image.size
+        image_width, image_height = control_image.size
     
-    image_after = resize_image(control_image, width, height, True)
+    image_after = resize_image(control_image, width, height, False)
     print(f"generate control image success: {image_width}x{image_height} => {width}x{height}")
 
     return image_after