up 3

app.py

@@ -3,9 +3,9 @@ import io
 from PIL import Image
 import numpy as np
 
-from config import WIDTH, HEIGHT
+import config
 from models import make_image_controlnet, make_inpainting
-from preprocessing import preprocess_seg_mask, get_image, get_mask
+from preprocessing import get_mask
 
 def image_to_byte_array(image: Image) -> bytes:
     # BytesIO is a fake file stored in memory
@@ -20,15 +20,19 @@ def predict(input_img1,
             input_img2,
             positive_prompt,
             negative_prompt,
-            num_of_images
+            num_of_images,
+            resolution
             ):
 
     print("predict")
     # input_img1 = Image.fromarray(input_img1)
     # input_img2 = Image.fromarray(input_img2)
 
-    input_img1 = input_img1.resize((WIDTH, HEIGHT))
-    input_img2 = input_img2.resize((WIDTH, WIDTH))
+    config.WIDTH = resolution
+    config.HEIGHT = resolution
+
+    input_img1 = input_img1.resize((config.WIDTH, config.HEIGHT))
+    input_img2 = input_img2.resize((config.WIDTH, config.HEIGHT))
 
     canvas_mask = np.array(input_img2)
     mask = get_mask(canvas_mask)
@@ -56,7 +60,8 @@ app = gr.Interface(
             gr.Image(label="mask", sources=['upload'], type="pil"),
             gr.Textbox(label="positive_prompt"),
             gr.Textbox(label="negative_prompt"),
-            gr.Number(label="num_of_images")
+            gr.Number(label="num_of_images"),
+            gr.Number(label="resolution")
             ],
     outputs= [
         gr.Image(label="resp0"),

models.py

@@ -69,7 +69,8 @@ def make_image_controlnet(image: np.ndarray,
 def make_inpainting(positive_prompt: str,
                     image: Image,
                     mask_image: np.ndarray,
-                    negative_prompt: str = "") -> List[Image.Image]:
+                    negative_prompt: str = "",
+                    rezolution:int=512) -> List[Image.Image]:
     """Method to make inpainting
     Args:
         positive_prompt (str): positive prompt string
@@ -90,8 +91,8 @@ def make_inpainting(positive_prompt: str,
                     prompt=positive_prompt,
                     negative_prompt=negative_prompt,
                     num_inference_steps=50,
-                    height=HEIGHT,
-                    width=WIDTH,
+                    height=rezolution,
+                    width=rezolution,
                     )
     print("RESP !!!!",resp)
     generated_image = resp.images[0]

preprocessing.py

@@ -6,7 +6,8 @@ import numpy as np
 from PIL import Image, ImageFilter
 import streamlit as st
 
-from config import COLOR_RGB, WIDTH, HEIGHT
+import config
+
 # from enhance_config import ENHANCE_SETTINGS
 
 LOGGING = logging.getLogger(__name__)
@@ -35,7 +36,7 @@ def preprocess_seg_mask(canvas_seg, real_seg: Image.Image = None) -> Tuple[np.nd
     unique_colors = [color for color in unique_colors if np.sum(
         np.all(image_seg == color, axis=-1)) > 100]
 
-    unique_colors_exact = [color for color in unique_colors if color in COLOR_RGB]
+    unique_colors_exact = [color for color in unique_colors if color in config.COLOR_RGB]
 
     if real_seg is not None:
         overlay_seg = np.array(real_seg)
@@ -74,9 +75,9 @@ def get_image() -> np.ndarray:
     if 'initial_image' in st.session_state and st.session_state['initial_image'] is not None:
         initial_image = st.session_state['initial_image']
         if isinstance(initial_image, Image.Image):
-            return np.array(initial_image.resize((WIDTH, HEIGHT)))
+            return np.array(initial_image.resize((config.WIDTH, config.HEIGHT)))
         else:
-            return np.array(Image.fromarray(initial_image).resize((WIDTH, HEIGHT)))
+            return np.array(Image.fromarray(initial_image).resize((config.WIDTH, config.HEIGHT)))
     else:
         return None