fix depth estimation

app.py

@@ -6,6 +6,7 @@ import diffuserslocal.src.diffusers as diffusers
 from share_btn import community_icon_html, loading_icon_html, share_js
 from diffuserslocal.src.diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_ldm3d_inpaint import StableDiffusionLDM3DInpaintPipeline
 from PIL import Image
+import numpy as np
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -33,7 +34,8 @@ else:
 
 
 def estimate_depth(image: Image) -> Image:
-
+    image= image.resize((384,384))
+    image = np.array(image)
     input_batch = transform(image).to(device)
 
     with torch.no_grad():
@@ -41,13 +43,18 @@ def estimate_depth(image: Image) -> Image:
 
         prediction = torch.nn.functional.interpolate(
             prediction.unsqueeze(1),
-            size=image.size,
+            size=image.shape[:2],
             mode="bicubic",
             align_corners=False,
         ).squeeze()
 
-    return Image.fromarray(prediction.cpu().numpy())
-     
+    output = prediction.cpu().numpy()
+
+    output= 255 * output/np.max(output)
+    
+    return Image.fromarray(output.astype("uint8"))
+
+
 def read_content(file_path: str) -> str:
     """read the content of target file
     """
@@ -60,6 +67,9 @@ def predict(dict, depth, prompt="", negative_prompt="", guidance_scale=7.5, step
     if negative_prompt == "":
         negative_prompt = None
     scheduler_class_name = scheduler.split("-")[0]
+
+    if (depth is None):
+        depth_image = estimate_depth(image)
     
     scheduler = getattr(diffusers, scheduler_class_name)
     pipe.scheduler = scheduler.from_pretrained("Intel/ldm3d-4c", subfolder="scheduler")
@@ -117,6 +127,8 @@ 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):
                         with gr.Row():
                             prompt = gr.Textbox(placeholder="Your prompt (what you want in place of what is erased)", show_label=False, elem_id="prompt")
@@ -140,10 +152,6 @@ with image_blocks as demo:
                         community_icon = gr.HTML(community_icon_html)
                         loading_icon = gr.HTML(loading_icon_html)
                         share_button = gr.Button("Share to community", elem_id="share-btn",visible=True)
-            
-
-    if (depth is None):
-         depth = estimate_depth(image)
 
     btn.click(fn=predict, inputs=[image, depth, prompt, negative_prompt, guidance_scale, steps, strength, scheduler], outputs=[image_out, depth_out, share_btn_container], api_name='run')
     prompt.submit(fn=predict, inputs=[image, depth, prompt, negative_prompt, guidance_scale, steps, strength, scheduler], outputs=[image_out, depth_out, share_btn_container])