Update

model.py

@@ -12,6 +12,7 @@ from diffusers import (ControlNetModel, DiffusionPipeline,
 
 from cv_utils import resize_image
 from preprocessor import Preprocessor
+from settings import MAX_IMAGE_RESOLUTION, MAX_NUM_IMAGES
 
 CONTROLNET_MODEL_IDS = {
     'Openpose': 'lllyasviel/control_v11p_sd15_openpose',
@@ -141,6 +142,11 @@ class Model:
         low_threshold: int,
         high_threshold: int,
     ) -> list[PIL.Image.Image]:
+        if image_resolution > MAX_IMAGE_RESOLUTION:
+            raise ValueError
+        if num_images > MAX_NUM_IMAGES:
+            raise ValueError
+
         self.preprocessor.load('Canny')
         control_image = self.preprocessor(image=image,
                                           low_threshold=low_threshold,
@@ -175,6 +181,11 @@ class Model:
         value_threshold: float,
         distance_threshold: float,
     ) -> list[PIL.Image.Image]:
+        if image_resolution > MAX_IMAGE_RESOLUTION:
+            raise ValueError
+        if num_images > MAX_NUM_IMAGES:
+            raise ValueError
+
         self.preprocessor.load('MLSD')
         control_image = self.preprocessor(
             image=image,
@@ -210,6 +221,11 @@ class Model:
         seed: int,
         preprocessor_name: str,
     ) -> list[PIL.Image.Image]:
+        if image_resolution > MAX_IMAGE_RESOLUTION:
+            raise ValueError
+        if num_images > MAX_NUM_IMAGES:
+            raise ValueError
+
         if preprocessor_name == 'None':
             image = HWC3(image)
             image = resize_image(image, resolution=image_resolution)
@@ -255,6 +271,11 @@ class Model:
         guidance_scale: float,
         seed: int,
     ) -> list[PIL.Image.Image]:
+        if image_resolution > MAX_IMAGE_RESOLUTION:
+            raise ValueError
+        if num_images > MAX_NUM_IMAGES:
+            raise ValueError
+
         image = image_and_mask['mask']
         image = HWC3(image)
         image = resize_image(image, resolution=image_resolution)
@@ -287,6 +308,11 @@ class Model:
         seed: int,
         preprocessor_name: str,
     ) -> list[PIL.Image.Image]:
+        if image_resolution > MAX_IMAGE_RESOLUTION:
+            raise ValueError
+        if num_images > MAX_NUM_IMAGES:
+            raise ValueError
+
         if preprocessor_name == 'None':
             image = HWC3(image)
             image = resize_image(image, resolution=image_resolution)
@@ -338,6 +364,11 @@ class Model:
         seed: int,
         preprocessor_name: str,
     ) -> list[PIL.Image.Image]:
+        if image_resolution > MAX_IMAGE_RESOLUTION:
+            raise ValueError
+        if num_images > MAX_NUM_IMAGES:
+            raise ValueError
+
         if preprocessor_name == 'None':
             image = HWC3(image)
             image = resize_image(image, resolution=image_resolution)
@@ -377,6 +408,11 @@ class Model:
         seed: int,
         preprocessor_name: str,
     ) -> list[PIL.Image.Image]:
+        if image_resolution > MAX_IMAGE_RESOLUTION:
+            raise ValueError
+        if num_images > MAX_NUM_IMAGES:
+            raise ValueError
+
         if preprocessor_name == 'None':
             image = HWC3(image)
             image = resize_image(image, resolution=image_resolution)
@@ -415,6 +451,11 @@ class Model:
         seed: int,
         preprocessor_name: str,
     ) -> list[PIL.Image.Image]:
+        if image_resolution > MAX_IMAGE_RESOLUTION:
+            raise ValueError
+        if num_images > MAX_NUM_IMAGES:
+            raise ValueError
+
         if preprocessor_name == 'None':
             image = HWC3(image)
             image = resize_image(image, resolution=image_resolution)
@@ -453,6 +494,11 @@ class Model:
         seed: int,
         preprocessor_name: str,
     ) -> list[PIL.Image.Image]:
+        if image_resolution > MAX_IMAGE_RESOLUTION:
+            raise ValueError
+        if num_images > MAX_NUM_IMAGES:
+            raise ValueError
+
         if preprocessor_name == 'None':
             image = HWC3(image)
             image = resize_image(image, resolution=image_resolution)
@@ -491,6 +537,11 @@ class Model:
         seed: int,
         preprocessor_name: str,
     ) -> list[PIL.Image.Image]:
+        if image_resolution > MAX_IMAGE_RESOLUTION:
+            raise ValueError
+        if num_images > MAX_NUM_IMAGES:
+            raise ValueError
+
         if preprocessor_name in ['None', 'None (anime)']:
             image = HWC3(image)
             image = resize_image(image, resolution=image_resolution)
@@ -540,6 +591,11 @@ class Model:
         seed: int,
         preprocessor_name: str,
     ) -> list[PIL.Image.Image]:
+        if image_resolution > MAX_IMAGE_RESOLUTION:
+            raise ValueError
+        if num_images > MAX_NUM_IMAGES:
+            raise ValueError
+
         if preprocessor_name == 'None':
             image = HWC3(image)
             image = resize_image(image, resolution=image_resolution)
@@ -575,6 +631,11 @@ class Model:
         guidance_scale: float,
         seed: int,
     ) -> list[PIL.Image.Image]:
+        if image_resolution > MAX_IMAGE_RESOLUTION:
+            raise ValueError
+        if num_images > MAX_NUM_IMAGES:
+            raise ValueError
+
         image = HWC3(image)
         image = resize_image(image, resolution=image_resolution)
         control_image = PIL.Image.fromarray(image)