neeew

app.py

@@ -51,11 +51,6 @@ def estimate_depth(image: Image) -> Image:
     output = prediction.cpu().numpy()
 
     output= 255 * output/np.max(output)
-
-    # If 3 channels convert to 1
-    if (len(output.shape) == 3):
-        if (output.shape[2] == 3):
-            output = output[:, :, 0]
     
     return Image.fromarray(output.astype("uint8"))
 
@@ -73,8 +68,10 @@ def predict(dict, depth, prompt="", negative_prompt="", guidance_scale=7.5, step
         negative_prompt = None
     scheduler_class_name = scheduler.split("-")[0]
 
-    if (depth is None):
+    if (depth.value == None):
         depth_image = estimate_depth(image)
+    else:
+        depth_image = depth.convert("L")
     
     scheduler = getattr(diffusers, scheduler_class_name)
     pipe.scheduler = scheduler.from_pretrained("Intel/ldm3d-4c", subfolder="scheduler")
@@ -132,6 +129,7 @@ with image_blocks as demo:
                 with gr.Column():
                     image = gr.Image(source='upload', tool='sketch', elem_id="image_upload", type="pil", label="Upload",height=400)
                     depth = gr.Image(source='upload', elem_id="depth_upload", type="pil", label="Upload",height=400)
+
                     print(depth)
 
                     with gr.Row(elem_id="prompt-container", mobile_collapse=False, equal_height=True):