Update app.py

app.py

@@ -16,8 +16,6 @@ import gradio as gr
 from huggingface_hub import hf_hub_download, snapshot_download
 from ip_adapter import IPAdapterXL
 from safetensors.torch import load_file
-from torchvision.models.detection import deeplabv3_resnet101
-from torchvision.transforms import functional as F
 
 snapshot_download(
     repo_id="h94/IP-Adapter", allow_patterns="sdxl_models/*", local_dir="."
@@ -138,43 +136,6 @@ def run_for_examples(style_image, source_image, prompt, scale, control_scale):
         neg_content_scale=0,
     )
 
-# Add the background removal function
-def remove_background(input_image):
-    # Load the deep learning model
-    model = deeplabv3_resnet101(pretrained=True)
-    model.eval()
-
-    # Preprocess the image
-    preprocess = transforms.Compose([
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-    ])
-    input_tensor = preprocess(input_image)
-    input_batch = input_tensor.unsqueeze(0)  # Create a mini-batch as expected by the model
-
-    # Move the input and model to GPU for speed if available
-    if torch.cuda.is_available():
-        input_batch = input_batch.to('cuda')
-        model.to('cuda')
-
-    with torch.no_grad():
-        output = model(input_batch)['out'][0]
-    output_predictions = output.argmax(0)
-
-    # Create a binary (black and white) mask of the profile foreground
-    mask = output_predictions.byte().cpu().numpy()
-    background = np.zeros(mask.shape)
-    bin_mask = np.where(mask, 255, background).astype(np.uint8)
-
-    # Create a transparent foreground
-    b, g, r = cv2.split(np.array(input_image).astype('uint8'))
-    a = np.ones(bin_mask.shape, dtype='uint8') * 255
-    alpha_im = cv2.merge([b, g, r, a], 4)
-    bg = np.zeros(alpha_im.shape)
-    new_mask = np.stack([bin_mask, bin_mask, bin_mask, bin_mask], axis=2)
-    foreground = np.where(new_mask, alpha_im, bg).astype(np.uint8)
-
-    return foreground
 
 @spaces.GPU
 def create_image(
@@ -221,17 +182,6 @@ def create_image(
         cv_input_image = pil_to_cv2(input_image)
         detected_map = cv2.Canny(cv_input_image, 50, 200)
         canny_map = Image.fromarray(cv2.cvtColor(detected_map, cv2.COLOR_BGR2RGB))
-        
-        # Remove background from the input image
-        foreground = remove_background(input_image)
-        # Convert the foreground back to a PIL image if necessary
-        foreground_pil = Image.fromarray(foreground)
-        
-        # Use foreground_pil instead of input_image for further processing
-        # Note: You might need to adjust the following lines based on how you intend to use the foreground_pil
-        # For example, if you're passing it to the IP-Adapter, ensure it's in the correct format
-        
-        # Continue with the existing logic for generating the image...
     else:
         canny_map = Image.new("RGB", (1024, 1024), color=(255, 255, 255))
         control_scale = 0
@@ -239,7 +189,7 @@ def create_image(
     if float(control_scale) == 0:
         canny_map = canny_map.resize((1024, 1024))
 
-    if len(neg_content_prompt) > 0 and neg_content_scale!= 0:
+    if len(neg_content_prompt) > 0 and neg_content_scale != 0:
         images = ip_model.generate(
             pil_image=image_pil,
             prompt=prompt,