Update utils/i2i.py

utils/i2i.py

@@ -20,6 +20,46 @@ def seed_everything(seed: int) -> torch.Generator:
     generator.manual_seed(seed)
     return generator
 
+# Function to properly handle input image transparency
+def prepare_transparent_image(input_image):
+    """
+    Ensures the input image has proper transparency.
+    Converts the image to RGBA if it's not already and enhances partial transparency.
+    """
+    # Convert to RGBA if not already
+    if input_image.mode != 'RGBA':
+        input_image = input_image.convert('RGBA')
+    
+    # Get image data as numpy array
+    img_array = np.array(input_image)
+    
+    # Print alpha channel stats for debugging
+    alpha = img_array[:, :, 3]
+    print(f"Input alpha channel stats - min: {alpha.min()}, max: {alpha.max()}, mean: {alpha.mean()}")
+    
+    # Check if the image has partial transparency (alpha values between 1-254)
+    if alpha.min() < 255 and np.any((alpha > 0) & (alpha < 255)):
+        print("Detected partial transparency, enhancing alpha channel...")
+        
+        # For partially transparent pixels (alpha < 200), make them fully transparent
+        partial_mask = alpha < 200
+        img_array[partial_mask, 3] = 0
+        
+        # For mostly opaque pixels (alpha >= 200), make them fully opaque
+        opaque_mask = alpha >= 200
+        img_array[opaque_mask, 3] = 255
+        
+        # Create new PIL image with enhanced alpha
+        enhanced_image = Image.fromarray(img_array, 'RGBA')
+        
+        # Print updated alpha stats
+        enhanced_alpha = np.array(enhanced_image)[:, :, 3]
+        print(f"Enhanced alpha channel stats - min: {enhanced_alpha.min()}, max: {enhanced_alpha.max()}, mean: {enhanced_alpha.mean()}")
+        
+        return enhanced_image
+    
+    return input_image
+
 # Initialize the pipeline
 i2i_pipe = FluxImg2ImgPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-dev",
@@ -70,6 +110,9 @@ def i2i_gen(
     gc.collect()
     
     try:
+        # Prepare the input image for proper transparency handling
+        input_image = prepare_transparent_image(input_image)
+        
         # Process the input image
         original_image = (transforms.ToTensor()(input_image)).unsqueeze(0)
         
@@ -222,7 +265,24 @@ def i2i_gen(
         # Convert to image - EXACTLY as in demo_i2i.py
         x = x.clamp(0, 1)
         x = x.permute(0, 2, 3, 1)
-        img = Image.fromarray((x*255).float().cpu().numpy().astype(np.uint8)[0])
+        img_array = (x*255).float().cpu().numpy().astype(np.uint8)[0]
+        
+        # Ensure the output image has proper transparency
+        if img_array.shape[2] == 4:
+            # Print alpha channel stats for debugging
+            alpha = img_array[:, :, 3]
+            print(f"Output alpha channel stats - min: {alpha.min()}, max: {alpha.max()}, mean: {alpha.mean()}")
+            
+            # Make sure partially transparent pixels are handled properly
+            # For partially transparent pixels (alpha < 200), make them fully transparent
+            partial_mask = alpha < 200
+            img_array[partial_mask, 3] = 0
+            
+            # For mostly opaque pixels (alpha >= 200), make them fully opaque
+            opaque_mask = alpha >= 200
+            img_array[opaque_mask, 3] = 255
+        
+        img = Image.fromarray(img_array, 'RGBA' if img_array.shape[2] == 4 else 'RGB')
         
         # Clean up
         del original_x, x
@@ -238,4 +298,4 @@ def i2i_gen(
         traceback.print_exc()
         torch.cuda.empty_cache()
         gc.collect()
-        return None
+        return None