update

.ipynb_checkpoints/app-checkpoint.py

@@ -41,10 +41,10 @@ def load_pth_data(pth_path):
 def infer_fn(prompt, mask_edited, raw_tensor_state, mask_tensor_state, cond_tensor_state):
     try:
         if isinstance(mask_edited, dict):
-            mask_edited = mask_edited["mask"]
+            mask_edited = mask_edited["composite"]
 
         mask_edited_tensor = transforms.ToTensor()(mask_edited)
-        mask_edited_tensor = mask_edited_tensor[:1].unsqueeze(0).half()
+        mask_edited_tensor = 1-mask_edited_tensor[:1].unsqueeze(0).half()
 
         raw_t = raw_tensor_state.half()
         cond_t = cond_tensor_state.half()
@@ -60,7 +60,7 @@ def infer_fn(prompt, mask_edited, raw_tensor_state, mask_tensor_state, cond_tens
             control_image=cond_t
         ).images[0]
 
-        return tensor_to_pil(result)
+        return tensor_to_pil(result), mask_edited
 
     except Exception as e:
         traceback.print_exc()
@@ -115,11 +115,11 @@ def build_demo():
                 mask_tensor_state,
                 cond_tensor_state
             ],
-            outputs=[output_image]
+            outputs=[output_image, rgb_display]
         )
 
     return demo
 
 if __name__ == "__main__":
     demo = build_demo()
-    demo.launch()
+    demo.launch(server_name="0.0.0.0", server_port=9112)

app.py

@@ -41,10 +41,10 @@ def load_pth_data(pth_path):
 def infer_fn(prompt, mask_edited, raw_tensor_state, mask_tensor_state, cond_tensor_state):
     try:
         if isinstance(mask_edited, dict):
-            mask_edited = mask_edited["mask"]
+            mask_edited = mask_edited["composite"]
 
         mask_edited_tensor = transforms.ToTensor()(mask_edited)
-        mask_edited_tensor = mask_edited_tensor[:1].unsqueeze(0).half()
+        mask_edited_tensor = 1-mask_edited_tensor[:1].unsqueeze(0).half()
 
         raw_t = raw_tensor_state.half()
         cond_t = cond_tensor_state.half()
@@ -60,7 +60,7 @@ def infer_fn(prompt, mask_edited, raw_tensor_state, mask_tensor_state, cond_tens
             control_image=cond_t
         ).images[0]
 
-        return tensor_to_pil(result)
+        return tensor_to_pil(result), mask_edited
 
     except Exception as e:
         traceback.print_exc()
@@ -115,11 +115,11 @@ def build_demo():
                 mask_tensor_state,
                 cond_tensor_state
             ],
-            outputs=[output_image]
+            outputs=[output_image, rgb_display]
         )
 
     return demo
 
 if __name__ == "__main__":
     demo = build_demo()
-    demo.launch()
+    demo.launch(server_name="0.0.0.0", server_port=9112)