Update app_merged.py

app_merged.py

@@ -599,8 +599,7 @@ def encode_prompt_pair(positive_prompt, negative_prompt):
 
     return c, uc
 
-# @spaces.GPU(duration=60)
-# @torch.inference_mode()
+
 @spaces.GPU(duration=60)
 @torch.inference_mode()
 def infer(
@@ -616,6 +615,8 @@ def infer(
     progress=gr.Progress(track_tqdm=True),
 ):
     #logging.info(f"Input image shape: {image.shape}, dtype: {image.dtype}")
+
+    clear_memory()
     
     # Convert input to PIL if needed
     if isinstance(image, np.ndarray):
@@ -698,35 +699,11 @@ 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
-#     if img.shape[-1] == 4:
-#         # Use white background for alpha composition
-#         alpha = img[..., 3:] / 255.0
-#         rgb = img[..., :3]
-#         white_bg = np.ones_like(rgb) * 255
-#         img = (rgb * alpha + white_bg * (1 - alpha)).astype(np.uint8)
-    
-#     H, W, C = img.shape
-#     assert C == 3
-#     k = (256.0 / float(H * W)) ** 0.5
-#     feed = resize_without_crop(img, int(64 * round(W * k)), int(64 * round(H * k)))
-#     feed = numpy2pytorch([feed]).to(device=device, dtype=torch.float32)
-#     alpha = rmbg(feed)[0][0]
-#     alpha = torch.nn.functional.interpolate(alpha, size=(H, W), mode="bilinear")
-#     alpha = alpha.movedim(1, -1)[0]
-#     alpha = alpha.detach().float().cpu().numpy().clip(0, 1)
-    
-#     # Create RGBA image
-#     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
 @torch.inference_mode()
 def run_rmbg(image):
+    clear_memory()
     image_size = image.size
     input_images = transform_image(image).unsqueeze(0).to("cuda")
     # Prediction
@@ -893,6 +870,7 @@ def process(input_fg, prompt, image_width, image_height, num_samples, seed, step
 def extract_foreground(image):
     if image is None:
         return None, gr.update(visible=True), gr.update(visible=True)
+    clear_memory()
     #logging.info(f"Input image shape: {image.shape}, dtype: {image.dtype}")
     #result, rgba = run_rmbg(image)
     result = run_rmbg(image)
@@ -1077,6 +1055,9 @@ def use_orientation(selected_image:gr.SelectData):
 def process_image(input_image, input_text):
     """Main processing function for the Gradio interface"""
 
+
+     clear_memory()
+
     if isinstance(input_image, Image.Image):
         input_image = np.array(input_image)
 
@@ -1443,6 +1424,8 @@ imageresize = NODE_CLASS_MAPPINGS["ImageResize+"]()
 @spaces.GPU
 def generate_image(prompt, structure_image, style_image, depth_strength=15, style_strength=0.5, progress=gr.Progress(track_tqdm=True)) -> str:
     """Main generation function that processes inputs and returns the path to the generated image."""
+
+    clear_memory()
     with torch.inference_mode():
         # Set up CLIP
         clip_switch = cr_clip_input_switch.switch(
@@ -1571,6 +1554,8 @@ def generate_image(prompt, structure_image, style_image, depth_strength=15, styl
             images=get_value_at_index(decoded, 0),
         )
         saved_path = f"output/{saved['ui']['images'][0]['filename']}"
+
+        clear_memory()
         return saved_path
 
 # Create Gradio interface