hi-pepe

worker_runpod.py

@@ -7,6 +7,10 @@ from PIL import Image
 from comfy.sd import load_checkpoint_guess_config
 import nodes
 
+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():
     model_patcher, clip, vae, clipvision = load_checkpoint_guess_config("/content/ComfyUI/models/checkpoints/model.safetensors", output_vae=True, output_clip=True, embedding_directory=None)
 
@@ -48,59 +52,39 @@ def generate(input):
     Image.fromarray(np.array(decoded*255, dtype=np.uint8)[0]).save("/content/output_image.png")
 
     result = "/content/output_image.png"
+    response = None
     try:
-        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')
+        source_id = values['source_id']
+        del values['source_id']
+        source_channel = values['source_channel']     
+        del values['source_channel']
         job_id = values['job_id']
         del values['job_id']
         default_filename = os.path.basename(result)
-        with open(result, "rb") as file:
-            files = {default_filename: file.read()}
-        payload = {"content": f"{json.dumps(values)} <@{discord_id}>"}
+        files = {default_filename: open(result, "rb").read()}
+        payload = {"content": f"{json.dumps(values)} <@{source_id}>"}
         response = requests.post(
-            f"https://discord.com/api/v9/channels/{discord_channel}/messages",
+            f"https://discord.com/api/v9/channels/{source_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:
-        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"}
+        print(f"An unexpected error occurred: {e}")
     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})