Update utils/i2i.py

utils/i2i.py

@@ -20,46 +20,6 @@ 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",
@@ -110,9 +70,6 @@ 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)
         
@@ -265,24 +222,7 @@ 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_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')
+        img = Image.fromarray((x*255).float().cpu().numpy().astype(np.uint8)[0])
         
         # Clean up
         del original_x, x
@@ -298,4 +238,4 @@ def i2i_gen(
         traceback.print_exc()
         torch.cuda.empty_cache()
         gc.collect()
-        return None
+        return None