neeew

app.py

@@ -69,7 +69,7 @@ def predict(dict, depth, prompt="", negative_prompt="", guidance_scale=7.5, step
         negative_prompt = None
     scheduler_class_name = scheduler.split("-")[0]
 
-    init_image = cv2.cvtColor(cv2.resize(dict["image"], (512, 512)), cv2.COLOR_BGR2RGB)
+    init_image = cv2.resize(dict["image"], (512, 512))
 
     mask = Image.fromarray(cv2.resize(dict["mask"], (512, 512))[:,:,0])
     mask.save("temp_mask.jpg")
@@ -93,7 +93,11 @@ def predict(dict, depth, prompt="", negative_prompt="", guidance_scale=7.5, step
     
     output = pipe(prompt = prompt, negative_prompt=negative_prompt, image=init_image, mask_image=mask, depth_image=depth_image, guidance_scale=guidance_scale, num_inference_steps=int(steps), strength=strength)
     
-    return output.rgb[0], output.depth[0], gr.update(visible=True)
+    output_depth_vis = (output.depth[0] - np.min(output.depth[0])) / (np.max(output.depth[0]) - np.min(output.depth[0])) * 255
+    output_depth_vis = output_depth_vis.astype("uint8")
+    output_depth = Image.fromarray(output_depth_vis)
+    
+    return output.rgb[0], output_depth, gr.update(visible=True)
 
 
 css = '''