updates

predict.py

@@ -22,25 +22,28 @@ class Predictor(BasePredictor):
         guidance_scale: float = Input(description="Guidance scale for the model", default=3.0, ge=0.0, le=14.0),
         number_of_images: int = Input(description="Number of images to generate", default=1, ge=1, le=4),
         number_of_steps: int = Input(description="Number of steps for the model", default=10, ge=1, le=50),
-        base_image: Path = Input(description="Base image for the model"),
-        base_image_strength: float = Input(description="Base image strength for the model", default=0.15, ge=0.0, le=1.0),
-        composition_image: Path = Input(description="Composition image for the model"),
-        composition_image_strength: float = Input(description="Composition image strength for the model", default=1.0, ge=0.0, le=1.0),
-        style_image: Path = Input(description="Style image for the model"),
-        style_image_strength: float = Input(description="Style image strength for the model", default=1.0, ge=0.0, le=1.0),
-        identity_image: Path = Input(description="Identity image for the model"),
-        identity_image_strength: float = Input(description="Identity image strength for the model", default=1.0, ge=0.0, le=1.0),
-        depth_image: Path = Input(description="Depth image for the model", default=None),
+        image_url: Path = Input(description="Base image for the model"),
+        image_strength: float = Input(description="Base image strength for the model", default=0.15, ge=0.0, le=1.0),
+        composition_image_url: Path = Input(description="Composition image for the model"),
+        composition_strength: float = Input(description="Composition image strength for the model", default=1.0, ge=0.0, le=1.0),
+        style_image_url: Path = Input(description="Style image for the model"),
+        style_strength: float = Input(description="Style image strength for the model", default=1.0, ge=0.0, le=1.0),
+        identity_image_url: Path = Input(description="Identity image for the model"),
+        identity_strength: float = Input(description="Identity image strength for the model", default=1.0, ge=0.0, le=1.0),
+        depth_image_url: Path = Input(description="Depth image for the model", default=None),
         depth_image_strength: float = Input(description="Depth image strength for the model, if not supplied the composition image will be used for depth", default=0.5, ge=0.0, le=1.0),
     ) -> List[Path]:
         """Run a single prediction on the model"""
 
-        base_image = Image.open(base_image)
-        composition_image = Image.open(composition_image)
-        style_image = Image.open(style_image)
-        identity_image = Image.open(identity_image)
-        if depth_image is not None:
-            depth_image = Image.open(depth_image)
+        base_image = Image.open(image_url)
+        composition_image = Image.open(composition_image_url)
+        style_image = Image.open(style_image_url)
+        identity_image = Image.open(identity_image_url)
+        if depth_image_url is not None:
+            depth_image = Image.open(depth_image_url)
+        else:
+            depth_image = None
+
         images = self.omni_zero.generate(
             seed=seed,
             prompt=prompt,
@@ -49,13 +52,13 @@ class Predictor(BasePredictor):
             number_of_images=number_of_images,
             number_of_steps=number_of_steps,
             base_image=base_image,
-            base_image_strength=base_image_strength,
+            base_image_strength=image_strength,
             composition_image=composition_image,
-            composition_image_strength=composition_image_strength,
+            composition_image_strength=composition_strength,
             style_image=style_image,
-            style_image_strength=style_image_strength,
+            style_image_strength=style_strength,
             identity_image=identity_image,
-            identity_image_strength=identity_image_strength,
+            identity_image_strength=identity_strength,
             depth_image=depth_image,
             depth_image_strength=depth_image_strength,
         )

requirements.txt

@@ -3,10 +3,11 @@ diffusers==0.29.0
 controlnet_aux==0.0.8
 huggingface_hub==0.23.2
 insightface==0.7.3
+albumentations==1.4.3
 numpy==1.26.2
 opencv_contrib_python==4.9.0.80
 opencv_python==4.9.0.80
-opencv_python_headless==4.7.0.72
+opencv_python_headless==4.9.0.80
 Pillow==10.1.0
 torch==2.0.0
 torchsde==0.2.6
@@ -15,3 +16,4 @@ onnxruntime-gpu
 hf_transfer
 gradio
 spaces
+pydantic<2.0.0