Spaces:
Runtime error
Runtime error
Update worker_runpod.py
Browse files- worker_runpod.py +43 -27
worker_runpod.py
CHANGED
@@ -1,3 +1,5 @@
|
|
|
|
|
|
1 |
import torch
|
2 |
import random
|
3 |
import comfy
|
@@ -13,12 +15,6 @@ import server
|
|
13 |
from nodes import load_custom_node
|
14 |
from math import ceil, floor
|
15 |
|
16 |
-
import os, json, requests, runpod
|
17 |
-
|
18 |
-
discord_token = os.getenv('com_camenduru_discord_token')
|
19 |
-
web_uri = os.getenv('com_camenduru_web_uri')
|
20 |
-
web_token = os.getenv('com_camenduru_web_token')
|
21 |
-
|
22 |
def download_file(url, save_dir='/content/ComfyUI/input'):
|
23 |
os.makedirs(save_dir, exist_ok=True)
|
24 |
file_name = url.split('/')[-1]
|
@@ -196,39 +192,59 @@ def generate(input):
|
|
196 |
Image.fromarray(np.array(blending_image*255, dtype=np.uint8)[0]).save("/content/ultralytics.png")
|
197 |
|
198 |
result = "/content/ultralytics.png"
|
199 |
-
response = None
|
200 |
try:
|
201 |
-
|
202 |
-
del values['
|
203 |
-
|
204 |
-
del values['
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
205 |
job_id = values['job_id']
|
206 |
del values['job_id']
|
207 |
default_filename = os.path.basename(result)
|
208 |
-
|
209 |
-
|
|
|
210 |
response = requests.post(
|
211 |
-
f"https://discord.com/api/v9/channels/{
|
212 |
data=payload,
|
213 |
-
headers={"
|
214 |
files=files
|
215 |
)
|
216 |
response.raise_for_status()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
217 |
except Exception as e:
|
218 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
219 |
finally:
|
220 |
if os.path.exists(result):
|
221 |
os.remove(result)
|
222 |
|
223 |
-
if response and response.status_code == 200:
|
224 |
-
try:
|
225 |
-
payload = {"jobId": job_id, "result": response.json()['attachments'][0]['url']}
|
226 |
-
requests.post(f"{web_uri}/api/notify", data=json.dumps(payload), headers={'Content-Type': 'application/json', "authorization": f"{web_token}"})
|
227 |
-
except Exception as e:
|
228 |
-
print(f"An unexpected error occurred: {e}")
|
229 |
-
finally:
|
230 |
-
return {"result": response.json()['attachments'][0]['url']}
|
231 |
-
else:
|
232 |
-
return {"result": "ERROR"}
|
233 |
-
|
234 |
runpod.serverless.start({"handler": generate})
|
|
|
1 |
+
import os, json, requests, runpod
|
2 |
+
|
3 |
import torch
|
4 |
import random
|
5 |
import comfy
|
|
|
15 |
from nodes import load_custom_node
|
16 |
from math import ceil, floor
|
17 |
|
|
|
|
|
|
|
|
|
|
|
|
|
18 |
def download_file(url, save_dir='/content/ComfyUI/input'):
|
19 |
os.makedirs(save_dir, exist_ok=True)
|
20 |
file_name = url.split('/')[-1]
|
|
|
192 |
Image.fromarray(np.array(blending_image*255, dtype=np.uint8)[0]).save("/content/ultralytics.png")
|
193 |
|
194 |
result = "/content/ultralytics.png"
|
|
|
195 |
try:
|
196 |
+
notify_uri = values['notify_uri']
|
197 |
+
del values['notify_uri']
|
198 |
+
notify_token = values['notify_token']
|
199 |
+
del values['notify_token']
|
200 |
+
discord_id = values['discord_id']
|
201 |
+
del values['discord_id']
|
202 |
+
if(discord_id == "discord_id"):
|
203 |
+
discord_id = os.getenv('com_camenduru_discord_id')
|
204 |
+
discord_channel = values['discord_channel']
|
205 |
+
del values['discord_channel']
|
206 |
+
if(discord_channel == "discord_channel"):
|
207 |
+
discord_channel = os.getenv('com_camenduru_discord_channel')
|
208 |
+
discord_token = values['discord_token']
|
209 |
+
del values['discord_token']
|
210 |
+
if(discord_token == "discord_token"):
|
211 |
+
discord_token = os.getenv('com_camenduru_discord_token')
|
212 |
job_id = values['job_id']
|
213 |
del values['job_id']
|
214 |
default_filename = os.path.basename(result)
|
215 |
+
with open(result, "rb") as file:
|
216 |
+
files = {default_filename: file.read()}
|
217 |
+
payload = {"content": f"{json.dumps(values)} <@{discord_id}>"}
|
218 |
response = requests.post(
|
219 |
+
f"https://discord.com/api/v9/channels/{discord_channel}/messages",
|
220 |
data=payload,
|
221 |
+
headers={"Authorization": f"Bot {discord_token}"},
|
222 |
files=files
|
223 |
)
|
224 |
response.raise_for_status()
|
225 |
+
result_url = response.json()['attachments'][0]['url']
|
226 |
+
notify_payload = {"jobId": job_id, "result": result_url, "status": "DONE"}
|
227 |
+
web_notify_uri = os.getenv('com_camenduru_web_notify_uri')
|
228 |
+
web_notify_token = os.getenv('com_camenduru_web_notify_token')
|
229 |
+
if(notify_uri == "notify_uri"):
|
230 |
+
requests.post(web_notify_uri, data=json.dumps(notify_payload), headers={'Content-Type': 'application/json', "Authorization": web_notify_token})
|
231 |
+
else:
|
232 |
+
requests.post(web_notify_uri, data=json.dumps(notify_payload), headers={'Content-Type': 'application/json', "Authorization": web_notify_token})
|
233 |
+
requests.post(notify_uri, data=json.dumps(notify_payload), headers={'Content-Type': 'application/json', "Authorization": notify_token})
|
234 |
+
return {"jobId": job_id, "result": result_url, "status": "DONE"}
|
235 |
except Exception as e:
|
236 |
+
error_payload = {"jobId": job_id, "status": "FAILED"}
|
237 |
+
try:
|
238 |
+
if(notify_uri == "notify_uri"):
|
239 |
+
requests.post(web_notify_uri, data=json.dumps(error_payload), headers={'Content-Type': 'application/json', "Authorization": web_notify_token})
|
240 |
+
else:
|
241 |
+
requests.post(web_notify_uri, data=json.dumps(error_payload), headers={'Content-Type': 'application/json', "Authorization": web_notify_token})
|
242 |
+
requests.post(notify_uri, data=json.dumps(error_payload), headers={'Content-Type': 'application/json', "Authorization": notify_token})
|
243 |
+
except:
|
244 |
+
pass
|
245 |
+
return {"jobId": job_id, "result": f"FAILED: {str(e)}", "status": "FAILED"}
|
246 |
finally:
|
247 |
if os.path.exists(result):
|
248 |
os.remove(result)
|
249 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
250 |
runpod.serverless.start({"handler": generate})
|