Add traceback

app.py

@@ -12,6 +12,7 @@ import spaces
 from gradio_imageslider import ImageSlider
 from torchvision.transforms.functional import to_pil_image, to_tensor
 from PIL import ImageFilter
+import traceback
 
 
 def preprocess_image(input_image, device):
@@ -73,49 +74,52 @@ if is_cpu_offload_enabled:
 
 @spaces.GPU
 def remove(gradio_image, rm_guidance_scale=9, num_inference_steps=50, seed=42, strength=0.8):
-    generator = torch.Generator('cuda').manual_seed(seed)
-    prompt = "" # Set prompt to null
-
-    source_image_pure = gradio_image["background"]
-    mask_image_pure = gradio_image["layers"][0]
-    source_image = preprocess_image(source_image_pure.convert('RGB'), device)
-    mask = preprocess_mask(mask_image_pure, device)
-    
-    START_STEP = 0 # AAS start step
-    END_STEP = int(strength * num_inference_steps) # AAS end step
-    LAYER = 34 # 0~23down,24~33mid,34~69up /AAS start layer 
-    END_LAYER = 70 # AAS end layer
-    ss_steps = 9 # similarity suppression steps
-    ss_scale = 0.3 # similarity suppression scale
-    
-    image = pipeline(
-        prompt=prompt,
-        image=source_image,
-        mask_image=mask,
-        height=1024,
-        width=1024,
-        AAS=True, # enable AAS
-        strength=strength, # inpainting strength
-        rm_guidance_scale=rm_guidance_scale, # removal guidance scale
-        ss_steps = ss_steps, # similarity suppression steps
-        ss_scale = ss_scale, # similarity suppression scale
-        AAS_start_step=START_STEP, # AAS start step
-        AAS_start_layer=LAYER, # AAS start layer 
-        AAS_end_layer=END_LAYER, # AAS end layer
-        num_inference_steps=num_inference_steps, # number of inference steps # AAS_end_step = int(strength*num_inference_steps)
-        generator=generator,
-        guidance_scale=1
-    ).images[0]
-    print('Inferece: DONE.')
-    
-    # pil_mask = to_pil_image(mask.squeeze(0))
-    # pil_mask_blurred = pil_mask.filter(ImageFilter.GaussianBlur(radius=15))
-    # mask_blurred = to_tensor(pil_mask_blurred).unsqueeze_(0).to(mask.device)
-    # mask_f = 1-(1 - mask) * (1 - mask_blurred)
-    
-    # image_1 = image.unsqueeze(0)
-    
-    return source_image, image
+    try:
+        generator = torch.Generator('cuda').manual_seed(seed)
+        prompt = "" # Set prompt to null
+
+        source_image_pure = gradio_image["background"]
+        mask_image_pure = gradio_image["layers"][0]
+        source_image = preprocess_image(source_image_pure.convert('RGB'), device)
+        mask = preprocess_mask(mask_image_pure, device)
+        
+        START_STEP = 0 # AAS start step
+        END_STEP = int(strength * num_inference_steps) # AAS end step
+        LAYER = 34 # 0~23down,24~33mid,34~69up /AAS start layer 
+        END_LAYER = 70 # AAS end layer
+        ss_steps = 9 # similarity suppression steps
+        ss_scale = 0.3 # similarity suppression scale
+        
+        image = pipeline(
+            prompt=prompt,
+            image=source_image,
+            mask_image=mask,
+            height=1024,
+            width=1024,
+            AAS=True, # enable AAS
+            strength=strength, # inpainting strength
+            rm_guidance_scale=rm_guidance_scale, # removal guidance scale
+            ss_steps = ss_steps, # similarity suppression steps
+            ss_scale = ss_scale, # similarity suppression scale
+            AAS_start_step=START_STEP, # AAS start step
+            AAS_start_layer=LAYER, # AAS start layer 
+            AAS_end_layer=END_LAYER, # AAS end layer
+            num_inference_steps=num_inference_steps, # number of inference steps # AAS_end_step = int(strength*num_inference_steps)
+            generator=generator,
+            guidance_scale=1
+        ).images[0]
+        print('Inferece: DONE.')
+        
+        # pil_mask = to_pil_image(mask.squeeze(0))
+        # pil_mask_blurred = pil_mask.filter(ImageFilter.GaussianBlur(radius=15))
+        # mask_blurred = to_tensor(pil_mask_blurred).unsqueeze_(0).to(mask.device)
+        # mask_f = 1-(1 - mask) * (1 - mask_blurred)
+        
+        # image_1 = image.unsqueeze(0)
+        
+        return source_image, image
+    except:
+        print(traceback.format_exc())
     
 
 title = """<h1 align="center">Object Remove</h1>"""