Update flow

app.py

@@ -8,10 +8,9 @@ from diffusers.utils import load_image
 from torchvision.transforms.functional import to_tensor, gaussian_blur
 from matplotlib import pyplot as plt
 import gradio as gr
-import spaces
 from gradio_imageslider import ImageSlider
 from torchvision.transforms.functional import to_pil_image, to_tensor
-from PIL import ImageFilter
+from PIL import ImageFilter, Image
 import traceback
 
 
@@ -22,19 +21,34 @@ def preprocess_image(input_image, device):
         image = image.expand(-1, 3, -1, -1)
     image = F.interpolate(image, (1024, 1024))
     image = image.to(dtype).to(device)
-    
+
     return image
 
 
 def preprocess_mask(input_mask, device):
-    mask = to_tensor(input_mask.convert('L'))
+    # Split the channels
+    r, g, b, alpha = input_mask.split()
+
+    # Create a new image where:
+    # - Black areas (where RGB = 0) become white (255).
+    # - Transparent areas (where alpha = 0) become black (0).
+    new_mask = Image.new("L", input_mask.size)
+
+    for x in range(input_mask.width):
+        for y in range(input_mask.height):
+            if alpha.getpixel((x, y)) == 0:  # Transparent pixel
+                new_mask.putpixel((x, y), 0)  # Set to black
+            else:  # Non-transparent pixel (originally black in the mask)
+                new_mask.putpixel((x, y), 255)  # Set to white
+    
+    mask = to_tensor(new_mask.convert('L'))
     mask = mask.unsqueeze_(0).float()  # 0 or 1
     mask = F.interpolate(mask, (1024, 1024))
     mask = gaussian_blur(mask, kernel_size=(77, 77))
     mask[mask < 0.1] = 0
     mask[mask >= 0.1] = 1
     mask = mask.to(dtype).to(device)
-    
+
     return mask
 
 
@@ -42,7 +56,7 @@ def make_redder(img, mask, increase_factor=0.4):
     img_redder = img.clone()
     mask_expanded = mask.expand_as(img)
     img_redder[0][mask_expanded[0] == 1] = torch.clamp(img_redder[0][mask_expanded[0] == 1] + increase_factor, 0, 1)
-    
+
     return img_redder
 
 
@@ -67,29 +81,28 @@ pipeline = DiffusionPipeline.from_pretrained(
 
 if is_attention_slicing_enabled:
     pipeline.enable_attention_slicing()
-    
+
 if is_cpu_offload_enabled:
     pipeline.enable_model_cpu_offload()
 
 
-@spaces.GPU
 def remove(gradio_image, rm_guidance_scale=9, num_inference_steps=50, seed=42, strength=0.8):
     try:
-        generator = torch.Generator(device).manual_seed(seed)
+        generator = torch.Generator('cuda').manual_seed(seed)
         prompt = "" # Set prompt to null
 
         source_image_pure = gradio_image["background"]
         mask_image_pure = gradio_image["layers"][0]
         source_image = preprocess_image(source_image_pure.convert('RGB'), device)
         mask = preprocess_mask(mask_image_pure, device)
-        
+
         START_STEP = 0 # AAS start step
         END_STEP = int(strength * num_inference_steps) # AAS end step
-        LAYER = 34 # 0~23down,24~33mid,34~69up /AAS start layer 
+        LAYER = 34 # 0~23down,24~33mid,34~69up /AAS start layer
         END_LAYER = 70 # AAS end layer
         ss_steps = 9 # similarity suppression steps
         ss_scale = 0.3 # similarity suppression scale
-        
+
         image = pipeline(
             prompt=prompt,
             image=source_image,
@@ -102,26 +115,25 @@ def remove(gradio_image, rm_guidance_scale=9, num_inference_steps=50, seed=42, s
             ss_steps = ss_steps, # similarity suppression steps
             ss_scale = ss_scale, # similarity suppression scale
             AAS_start_step=START_STEP, # AAS start step
-            AAS_start_layer=LAYER, # AAS start layer 
+            AAS_start_layer=LAYER, # AAS start layer
             AAS_end_layer=END_LAYER, # AAS end layer
             num_inference_steps=num_inference_steps, # number of inference steps # AAS_end_step = int(strength*num_inference_steps)
             generator=generator,
-            guidance_scale=1,
-            output_type='pt'
+            guidance_scale=1
         ).images[0]
         print('Inferece: DONE.')
-        
+
         pil_mask = to_pil_image(mask.squeeze(0))
         pil_mask_blurred = pil_mask.filter(ImageFilter.GaussianBlur(radius=15))
         mask_blurred = to_tensor(pil_mask_blurred).unsqueeze_(0).to(mask.device)
         mask_f = 1-(1 - mask) * (1 - mask_blurred)
-        
-        image_1 = image.unsqueeze(0)
-        
-        return source_image_pure, pil_mask, image_1
+
+        # image_1 = image.unsqueeze(0)
+
+        return source_image_pure, pil_mask, image
     except:
         print(traceback.format_exc())
-    
+
 
 title = """<h1 align="center">Object Remove</h1>"""
 with gr.Blocks() as demo:
@@ -157,7 +169,7 @@ with gr.Blocks() as demo:
                     step=0.1,
                     label="Strength"
                 )
-            
+
             input_image = gr.ImageMask(
                 type="pil", label="Input Image",crop_size=(1200,1200), layers=False
             )
@@ -167,11 +179,11 @@ with gr.Blocks() as demo:
                     run_button = gr.Button("Generate")
 
             result = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery", columns=[3], rows=[1], object_fit="contain", height="auto")
-    
+
     run_button.click(
         fn=remove,
         inputs=[input_image, guidance_scale, num_steps, seed, strength],
         outputs=result,
     )
-    
-demo.queue(max_size=12).launch(share=False)
+
+demo.queue(max_size=12).launch(share=True)