Update app.py

app.py

@@ -1,7 +1,5 @@
 import io
 import os
-os.system("pip uninstall -y gradio")  
-os.system("pip install gradio==3.49.0")
 import shutil
 import uuid
 import torch
@@ -106,69 +104,77 @@ def predict(
     scale,
     image_paths,
     mask_paths
+
 ):
     global image_path, mask_path
     gr.Info(str(f"Set seed = {seed}"))
     if image_paths is not None:
-        input_image["image"] = load_image(image_paths).convert("RGB")
-        input_image["mask"] = load_image(mask_paths).convert("RGB")
+        input_image["background"] = load_image(image_paths).convert("RGB")
+        input_image["layers"][0] = load_image(mask_paths).convert("RGB")
         
-    size1, size2 = input_image["image"].convert("RGB").size
-
-    icc_profile = input_image["image"].info.get('icc_profile')
+    size1, size2 = input_image["background"].convert("RGB").size
+    icc_profile = input_image["background"].info.get('icc_profile')
     if icc_profile:
         gr.Info(str(f"Image detected to contain ICC profile, converting color space to sRGB..."))
         srgb_profile = ImageCms.createProfile("sRGB")
         io_handle = io.BytesIO(icc_profile)   
         src_profile = ImageCms.ImageCmsProfile(io_handle)  
-        input_image["image"] = ImageCms.profileToProfile(input_image["image"], src_profile, srgb_profile)
-        input_image["image"].info.pop('icc_profile', None)
+        input_image["background"] = ImageCms.profileToProfile(input_image["background"], src_profile, srgb_profile)
+        input_image["background"].info.pop('icc_profile', None)
 
     if size1 < size2:
-        input_image["image"] = input_image["image"].convert("RGB").resize((1024, int(size2 / size1 * 1024)))
+        input_image["background"] = input_image["background"].convert("RGB").resize((1024, int(size2 / size1 * 1024)))
     else:
-        input_image["image"] = input_image["image"].convert("RGB").resize((int(size1 / size2 * 1024), 1024))
+        input_image["background"] = input_image["background"].convert("RGB").resize((int(size1 / size2 * 1024), 1024))
 
-    img = np.array(input_image["image"].convert("RGB"))
+    img = np.array(input_image["background"].convert("RGB"))
 
     W = int(np.shape(img)[0] - np.shape(img)[0] % 8)
     H = int(np.shape(img)[1] - np.shape(img)[1] % 8)
 
-    input_image["image"] = input_image["image"].resize((H, W))
-    input_image["mask"] = input_image["mask"].resize((H, W))
+    input_image["background"] = input_image["background"].resize((H, W))
+    input_image["layers"][0] = input_image["layers"][0].resize((H, W))
 
     if seed == -1:
         seed = random.randint(1, 2147483647)
         set_seed(random.randint(1, 2147483647))
     else:
         set_seed(seed)
-
-
+    if image_paths is None:
+        img=input_image["layers"][0]
+        img_data = np.array(img)
+        alpha_channel = img_data[:, :, 3]
+        white_background = np.ones_like(alpha_channel) * 255
+        gray_image = white_background.copy()
+        gray_image[alpha_channel == 0] = 0
+        gray_image_pil = Image.fromarray(gray_image).convert('L')
+    else:
+        gray_image_pil = input_image["layers"][0]
     result = pipe(
         prompt=prompt,
-        control_image=input_image["image"].convert("RGB"),
-        control_mask=input_image["mask"].convert("RGB"),
+        control_image=input_image["background"].convert("RGB"),
+        control_mask=gray_image_pil.convert("RGB"),
         width=H,
         height=W,
         num_inference_steps=ddim_steps,
-        generator=torch.Generator(device).manual_seed(seed),
+        generator=torch.Generator("cuda").manual_seed(seed),
         guidance_scale=scale,
         max_sequence_length=512,
     ).images[0]
 
-    mask_np = np.array(input_image["mask"].convert("RGB"))
-    red = np.array(input_image["image"]).astype("float") * 1
+    mask_np = np.array(input_image["layers"][0].convert("RGB"))
+    red = np.array(input_image["background"]).astype("float") * 1
     red[:, :, 0] = 180.0
     red[:, :, 2] = 0
     red[:, :, 1] = 0
-    result_m = np.array(input_image["image"])
+    result_m = np.array(input_image["background"])
     result_m = Image.fromarray(
         (
             result_m.astype("float") * (1 - mask_np.astype("float") / 512.0) + mask_np.astype("float") / 512.0 * red
         ).astype("uint8")
     )
 
-    dict_res = [input_image["image"], input_image["mask"], result_m, result]
+    dict_res = [input_image["background"], input_image["layers"][0], result_m, result]
 
     dict_out = [result]
     image_path = None
@@ -182,6 +188,7 @@ def infer(
     seed,
     scale,
     removal_prompt,
+
 ):
     img_path = image_path
     msk_path = mask_path
@@ -205,7 +212,9 @@ def process_example(image_paths, mask_paths):
     mask_path = mask_paths
     return masked_image
 custom_css = """
+
 .contain { max-width: 1200px !important; }
+
 .custom-image {
     border: 2px dashed #7e22ce !important;
     border-radius: 12px !important;
@@ -215,6 +224,7 @@ custom_css = """
     border-color: #9333ea !important;
     box-shadow: 0 4px 15px rgba(158, 109, 202, 0.2) !important;
 }
+
 .btn-primary {
     background: linear-gradient(45deg, #7e22ce, #9333ea) !important;
     border: none !important;
@@ -227,14 +237,17 @@ custom_css = """
     padding: 16px !important;
     margin-top: 8px !important;
 }
+
 #inline-examples .thumbnail {
     border-radius: 8px !important;
     transition: transform 0.2s ease !important;
 }
+
 #inline-examples .thumbnail:hover {
     transform: scale(1.05);
     box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1);
 }
+
 .example-title h3 {
     margin: 0 0 12px 0 !important;
     color: #475569 !important;
@@ -242,11 +255,16 @@ custom_css = """
     display: flex !important;
     align-items: center !important;
 }
+
 .example-title h3::before {
     content: "📚";
     margin-right: 8px;
     font-size: 1.2em;
 }
+
+.row { align-items: stretch !important; }
+
+.panel { height: 100%; }
 """
 
 with gr.Blocks(
@@ -273,16 +291,15 @@ with gr.Blocks(
     </div>
     """)
 
-    with gr.Row(equal_height=True):
+    with gr.Row(equal_height=False):
         with gr.Column(scale=1, variant="panel"):
             gr.Markdown("## 📥 Input Panel")
             
             with gr.Group():
-                input_image = gr.Image(
+                input_image = gr.Sketchpad(
+                    sources=["upload"],
                     type="pil",
-                    tool="sketch",
                     label="Upload & Annotate",
-                    height=400,
                     elem_id="custom-image",
                     interactive=True
                 )