Update app.py

app.py

@@ -3,9 +3,9 @@ import numpy as np
 import random
 import torch
 import spaces
-
-
-
+from transformers import AutoModelForImageSegmentation
+from torchvision import transforms
+from typing import Union, Tuple
 from PIL import Image
 from diffusers import FlowMatchEulerDiscreteScheduler,DiffusionPipeline
 from optimization import optimize_pipeline_
@@ -17,7 +17,7 @@ from huggingface_hub import InferenceClient
 import math
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
-
+from typing import Union, Tuple
 from basicsr.archs.rrdbnet_arch import RRDBNet
 from basicsr.utils.download_util import load_file_from_url
 from realesrgan import RealESRGANer
@@ -37,6 +37,42 @@ import tempfile
 from PIL import Image
 import gradio as gr
 
+torch.set_float32_matmul_precision(["high", "highest"][0])
+
+birefnet = AutoModelForImageSegmentation.from_pretrained(
+    "ZhengPeng7/BiRefNet", trust_remote_code=True
+)
+birefnet.to("cuda")
+
+transform_image = transforms.Compose(
+    [
+        transforms.Resize((1024, 1024)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+    ]
+)
+
+@spaces.GPU
+def processRemove(image: Image.Image) -> Image.Image:
+    """
+    Apply BiRefNet-based image segmentation to remove the background.
+    This function preprocesses the input image, runs it through a BiRefNet segmentation model to obtain a mask,
+    and applies the mask as an alpha (transparency) channel to the original image.
+    Args:
+        image (PIL.Image): The input RGB image.
+    Returns:
+        PIL.Image: The image with the background removed, using the segmentation mask as transparency.
+    """
+    image_size = image.size
+    input_images = transform_image(image).unsqueeze(0).to("cuda")
+    # Prediction
+    with torch.no_grad():
+        preds = birefnet(input_images)[-1].sigmoid().cpu()
+    pred = preds[0].squeeze()
+    pred_pil = transforms.ToPILImage()(pred)
+    mask = pred_pil.resize(image_size)
+    image.putalpha(mask)
+    return image
 
 # --- Upscaling ---
 MAX_SEED = np.iinfo(np.int32).max
@@ -348,11 +384,17 @@ def infer(
         true_cfg_scale=true_guidance_scale,
         num_images_per_prompt=num_images_per_prompt,
     ).images
+
+    im = load_img(image[0], output_type="pil")
+    im = im.convert("RGB")
+    origin = im.copy()
+    removed_image = process(im)
+    
     if return_upscaled:
         upscaled = realesrgan(image[0], "realesr-general-x4v3", 0.5, 2)
         return image, upscaled, seed
     else:
-        return image, None, seed
+        return image, removed_image, seed
 
 
 # --- Examples and UI Layout ---