matte

matte_processor.py

@@ -4,6 +4,7 @@ import torch
 from transformers import VitMatteImageProcessor, VitMatteForImageMatting
 import math
 from pathlib import Path
+import numpy as np
 
 class VITMatteModel:
     def __init__(self, model, processor):
@@ -11,11 +12,14 @@ class VITMatteModel:
         self.processor = processor
 
 def load_VITMatte_model(local_files_only=False):
-    model = VitMatteForImageMatting.from_pretrained("hustvl/vitmatte-small-composition-1k", local_files_only=local_files_only)
+    model = VitMatteForImageMatting.from_pretrained("hustvl/vitmatte-small-composition-1k", local_files_only=local_files_only).to("cuda")
     processor = VitMatteImageProcessor.from_pretrained("hustvl/vitmatte-small-composition-1k", local_files_only=local_files_only)
     return VITMatteModel(model, processor)
 
-def generate_VITMatte(image, trimap, local_files_only=False, device="cpu", max_megapixels=2.0):
+# Load model globally
+vit_matte_model = load_VITMatte_model(local_files_only=False)
+
+def generate_VITMatte(image, trimap, local_files_only=False, max_megapixels=2.0):
     if image is None or trimap is None:
         return None
         
@@ -45,19 +49,11 @@ def generate_VITMatte(image, trimap, local_files_only=False, device="cpu", max_m
         trimap = trimap.resize((target_width, target_height), Image.BILINEAR)
         resized = True
 
-    # Set device
-    if device == "cuda" and not torch.cuda.is_available():
-        device = "cpu"
-    device = torch.device(device)
-
-    # Load and process
-    vit_matte_model = load_VITMatte_model(local_files_only=local_files_only)
-    vit_matte_model.model.to(device)
-
+    # Use global model
     inputs = vit_matte_model.processor(images=image, trimaps=trimap, return_tensors="pt")
     
     with torch.no_grad():
-        inputs = {k: v.to(device) for k, v in inputs.items()}
+        inputs = {k: v.to("cuda") for k, v in inputs.items()}
         predictions = vit_matte_model.model(**inputs).alphas
 
     if torch.cuda.is_available():
@@ -83,12 +79,6 @@ def create_matte_tab():
                 input_image = gr.Image(label="Input Image", type="numpy", height=256)
                 trimap_image = gr.Image(label="Trimap Image", type="numpy", height=256)
                 
-                device = gr.Radio(
-                    choices=["cpu", "cuda"], 
-                    value="cpu",
-                    label="Device"
-                )
-                
                 max_megapixels = gr.Slider(
                     minimum=0.5,
                     maximum=8.0,
@@ -113,7 +103,6 @@ def create_matte_tab():
                 input_image,
                 trimap_image,
                 local_files,
-                device,
                 max_megapixels
             ],
             outputs=output_image

requirements.txt

@@ -1,3 +1,4 @@
+transformers
 pixeloe
 torch
 torchvision