Update worker_runpod.py

worker_runpod.py

@@ -83,7 +83,10 @@ def generate(input):
         }
         supported_tasks[args.task](args).run()
 
-    result = "/content/result/diffbir_tost.png"
+    if args.version == "fr_bg":
+        result = "/content/result/restored_backgrounds/diffbir_tost.png"
+    else:
+        result = "/content/result/diffbir_tost.png"
     try:
         notify_uri = values['notify_uri']
         del values['notify_uri']
@@ -136,10 +139,6 @@ def generate(input):
             pass
         return {"jobId": job_id, "result": f"FAILED: {str(e)}", "status": "FAILED"}
     finally:
-        if os.path.exists(result):
-            os.remove(result)
-        if os.path.exists(input_image):
-            os.remove(input_image)
         if os.path.exists('/content/input'):
             shutil.rmtree('/content/input')
         if os.path.exists('/content/result'):