Update worker_runpod.py

worker_runpod.py

@@ -17,10 +17,6 @@ from transport import Sampler, create_transport
 from diffusers.models import AutoencoderKL
 from transformers import AutoModel, AutoTokenizer
 
-discord_token = os.getenv('com_camenduru_discord_token')
-web_uri = os.getenv('com_camenduru_web_uri')
-web_token = os.getenv('com_camenduru_web_token')
-
 with torch.inference_mode():
     path_type = "Linear" # ["Linear", "GVP", "VP"]
     prediction = "velocity" # ["velocity", "score", "noise"]
@@ -197,39 +193,59 @@ def generate(input):
             print(traceback.format_exc())
 
     result = "/content/out.png"
-    response = None
     try:
-        source_id = values['source_id']
-        del values['source_id']
-        source_channel = values['source_channel']     
-        del values['source_channel']
+        notify_uri = values['notify_uri']
+        del values['notify_uri']
+        notify_token = values['notify_token']
+        del values['notify_token']
+        discord_id = values['discord_id']
+        del values['discord_id']
+        if(discord_id == "discord_id"):
+            discord_id = os.getenv('com_camenduru_discord_id')
+        discord_channel = values['discord_channel']
+        del values['discord_channel']
+        if(discord_channel == "discord_channel"):
+            discord_channel = os.getenv('com_camenduru_discord_channel')
+        discord_token = values['discord_token']
+        del values['discord_token']
+        if(discord_token == "discord_token"):
+            discord_token = os.getenv('com_camenduru_discord_token')
         job_id = values['job_id']
         del values['job_id']
         default_filename = os.path.basename(result)
-        files = {default_filename: open(result, "rb").read()}
-        payload = {"content": f"{json.dumps(values)} <@{source_id}>"}
+        with open(result, "rb") as file:
+            files = {default_filename: file.read()}
+        payload = {"content": f"{json.dumps(values)} <@{discord_id}>"}
         response = requests.post(
-            f"https://discord.com/api/v9/channels/{source_channel}/messages",
+            f"https://discord.com/api/v9/channels/{discord_channel}/messages",
             data=payload,
-            headers={"authorization": f"Bot {discord_token}"},
+            headers={"Authorization": f"Bot {discord_token}"},
             files=files
         )
         response.raise_for_status()
+        result_url = response.json()['attachments'][0]['url']
+        notify_payload = {"jobId": job_id, "result": result_url, "status": "DONE"}
+        web_notify_uri = os.getenv('com_camenduru_web_notify_uri')
+        web_notify_token = os.getenv('com_camenduru_web_notify_token')
+        if(notify_uri == "notify_uri"):
+            requests.post(web_notify_uri, data=json.dumps(notify_payload), headers={'Content-Type': 'application/json', "Authorization": web_notify_token})
+        else:
+            requests.post(web_notify_uri, data=json.dumps(notify_payload), headers={'Content-Type': 'application/json', "Authorization": web_notify_token})
+            requests.post(notify_uri, data=json.dumps(notify_payload), headers={'Content-Type': 'application/json', "Authorization": notify_token})
+        return {"jobId": job_id, "result": result_url, "status": "DONE"}
     except Exception as e:
-        print(f"An unexpected error occurred: {e}")
+        error_payload = {"jobId": job_id, "status": "FAILED"}
+        try:
+            if(notify_uri == "notify_uri"):
+                requests.post(web_notify_uri, data=json.dumps(error_payload), headers={'Content-Type': 'application/json', "Authorization": web_notify_token})
+            else:
+                requests.post(web_notify_uri, data=json.dumps(error_payload), headers={'Content-Type': 'application/json', "Authorization": web_notify_token})
+                requests.post(notify_uri, data=json.dumps(error_payload), headers={'Content-Type': 'application/json', "Authorization": notify_token})
+        except:
+            pass
+        return {"jobId": job_id, "result": f"FAILED: {str(e)}", "status": "FAILED"}
     finally:
         if os.path.exists(result):
             os.remove(result)
 
-    if response and response.status_code == 200:
-        try:
-            payload = {"jobId": job_id, "result": response.json()['attachments'][0]['url']}
-            requests.post(f"{web_uri}/api/notify", data=json.dumps(payload), headers={'Content-Type': 'application/json', "authorization": f"{web_token}"})
-        except Exception as e:
-            print(f"An unexpected error occurred: {e}")
-        finally:
-            return {"result": response.json()['attachments'][0]['url']}
-    else:
-        return {"result": "ERROR"}
-
 runpod.serverless.start({"handler": generate})