mask fix

app.py

@@ -13,11 +13,11 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Inpainting pipeline
 
-
+'''
 unet = UNet2DConditionModel.from_pretrained("pablodawson/ldm3d-inpainting", cache_dir="cache", subfolder="unet", in_channels=9, low_cpu_mem_usage=False, ignore_mismatched_sizes=True)
 pipe = StableDiffusionLDM3DInpaintPipeline.from_pretrained("Intel/ldm3d-4c", cache_dir="cache" ).to(device)
 
-
+'''
 # Depth estimation
 model_type = "DPT_Large"     # MiDaS v3 - Large     (highest accuracy, slowest inference speed)
 #model_type = "DPT_Hybrid"   # MiDaS v3 - Hybrid    (medium accuracy, medium inference speed)
@@ -71,7 +71,8 @@ def predict(dict, depth, prompt="", negative_prompt="", guidance_scale=7.5, step
     scheduler_class_name = scheduler.split("-")[0]
 
     init_image = cv2.resize(dict["image"], (512, 512))
-    
+    mask = Image.fromarray(cv2.resize(dict["mask"], (512, 512))[:,:,0])
+
     if (depth is None):
         depth_image = estimate_depth(init_image)
     else:
@@ -81,7 +82,7 @@ def predict(dict, depth, prompt="", negative_prompt="", guidance_scale=7.5, step
     scheduler = getattr(diffusers, scheduler_class_name)
     pipe.scheduler = scheduler.from_pretrained("Intel/ldm3d-4c", subfolder="scheduler")
     
-    mask = cv2.resize(dict["mask"], (512, 512))
+    
     depth_image = depth_image.resize((512, 512))
     
     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)