Update app.py

app.py

@@ -120,6 +120,7 @@ unet = UNet2DConditionModel.from_pretrained(sd15_name, subfolder="unet")
 # Load model directly
 from transformers import AutoModelForImageSegmentation
 rmbg = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-1.4", trust_remote_code=True)
+rmbg = rmbg.to(device=device, dtype=torch.float32)  # Keep this as float32
 
 model = DepthAnythingV2(encoder='vits', features=64, out_channels=[48, 96, 192, 384])
 model.load_state_dict(torch.load('checkpoints/depth_anything_v2_vits.pth', map_location=device))
@@ -128,6 +129,7 @@ model.eval()
 
 # Change UNet
 
+
 with torch.no_grad():
     new_conv_in = torch.nn.Conv2d(8, unet.conv_in.out_channels, unet.conv_in.kernel_size, unet.conv_in.stride, unet.conv_in.padding)
     new_conv_in.weight.zero_()
@@ -314,7 +316,7 @@ def encode_prompt_pair(positive_prompt, negative_prompt):
 
     return c, uc
 
-
+@spaces.GPU(duration=60)
 @torch.inference_mode()
 def pytorch2numpy(imgs, quant=True):
     results = []
@@ -331,7 +333,7 @@ def pytorch2numpy(imgs, quant=True):
         results.append(y)
     return results
 
-
+@spaces.GPU(duration=60)
 @torch.inference_mode()
 def numpy2pytorch(imgs):
     h = torch.from_numpy(np.stack(imgs, axis=0)).float() / 127.0 - 1.0  # so that 127 must be strictly 0.0
@@ -359,7 +361,7 @@ def resize_without_crop(image, target_width, target_height):
     resized_image = pil_image.resize((target_width, target_height), Image.LANCZOS)
     return np.array(resized_image)
 
-
+@spaces.GPU(duration=60)
 @torch.inference_mode()
 def run_rmbg(img, sigma=0.0):
     # Convert RGBA to RGB if needed
@@ -384,6 +386,8 @@ def run_rmbg(img, sigma=0.0):
     rgba = np.dstack((img, alpha * 255)).astype(np.uint8)
     result = 127 + (img.astype(np.float32) - 127 + sigma) * alpha
     return result.clip(0, 255).astype(np.uint8), rgba
+
+@spaces.GPU(duration=60)
 @torch.inference_mode()
 def process(input_fg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, lowres_denoise, bg_source):
     clear_memory()