140 lines
4.6 KiB
Python
140 lines
4.6 KiB
Python
from comfy.model_management import InterruptProcessingException
|
|
from requests.adapters import HTTPAdapter
|
|
from server import PromptServer
|
|
from dotenv import load_dotenv
|
|
from pathlib import Path
|
|
from PIL import Image
|
|
import numpy as np
|
|
import requests
|
|
import random
|
|
import base64
|
|
import torch
|
|
import uuid
|
|
import os
|
|
import io
|
|
|
|
load_dotenv()
|
|
|
|
SESSION_TIMEOUT = 60
|
|
RUNWARE_API_BASE_URL = "https://api.runware.ai/v1"
|
|
session = requests.Session()
|
|
adapter = HTTPAdapter(pool_connections=10, pool_maxsize=10)
|
|
session.mount("http://", adapter)
|
|
session.mount("https://", adapter)
|
|
|
|
def getAPIKey():
|
|
apiKey = os.getenv("RUNWARE_API_KEY")
|
|
if apiKey and isinstance(apiKey, str) and len(apiKey) > 24:
|
|
return apiKey
|
|
return False
|
|
|
|
RUNWARE_API_KEY = getAPIKey()
|
|
|
|
def setAPIKey(apiKey: str):
|
|
global RUNWARE_API_KEY
|
|
comfyNodeRoot = Path(__file__).parent.parent.parent
|
|
envFilePath = comfyNodeRoot / '.env'
|
|
if not envFilePath.exists():
|
|
envFilePath.touch()
|
|
with open(envFilePath, 'r') as f:
|
|
lines = f.readlines()
|
|
key_exists = False
|
|
new_lines = []
|
|
for line in lines:
|
|
if line.startswith('RUNWARE_API_KEY='):
|
|
key_exists = True
|
|
new_lines.append(f'RUNWARE_API_KEY={apiKey}\n')
|
|
else:
|
|
new_lines.append(line)
|
|
if not key_exists:
|
|
new_lines.append(f'RUNWARE_API_KEY={apiKey}\n')
|
|
with open(envFilePath, 'w') as f:
|
|
f.writelines(new_lines)
|
|
RUNWARE_API_KEY = apiKey
|
|
os.environ["RUNWARE_API_KEY"] = apiKey
|
|
return True
|
|
|
|
def genRandSeed(minSeed = 1000, maxSeed = 9223372036854776000):
|
|
return random.randint(minSeed, maxSeed)
|
|
|
|
def genRandUUID():
|
|
return str(uuid.uuid4())
|
|
|
|
def checkAPIKey(apiKey):
|
|
headers = {
|
|
"Accept": "application/json",
|
|
"Content-Type": "application/json",
|
|
"Accept-Encoding": "gzip, deflate, br, zstd",
|
|
}
|
|
|
|
genConfig = [{
|
|
"taskType": "authentication",
|
|
"apiKey": apiKey,
|
|
}]
|
|
|
|
try:
|
|
genResult = session.post(RUNWARE_API_BASE_URL, headers=headers, json=genConfig, timeout=SESSION_TIMEOUT, allow_redirects=False, stream=True)
|
|
genResult = genResult.json()
|
|
if("errors" in genResult):
|
|
return str(genResult["errors"][0]["message"])
|
|
else:
|
|
return True
|
|
except Exception as e:
|
|
return False
|
|
|
|
def inferenecRequest(genConfig):
|
|
global RUNWARE_API_KEY
|
|
RUNWARE_API_KEY = os.getenv("RUNWARE_API_KEY")
|
|
headers = {
|
|
"Authorization": f"Bearer {RUNWARE_API_KEY}",
|
|
"Content-Type": "application/json",
|
|
"Accept-Encoding": "gzip, deflate, br, zstd",
|
|
}
|
|
|
|
try:
|
|
genResult = session.post(RUNWARE_API_BASE_URL, headers=headers, json=genConfig, timeout=SESSION_TIMEOUT, allow_redirects=False, stream=True)
|
|
genResult = genResult.json()
|
|
if("errors" in genResult):
|
|
raise Exception(genResult["errors"][0]["message"])
|
|
else:
|
|
return genResult
|
|
except requests.exceptions.Timeout:
|
|
raise Exception(f"Error: Request Timed Out After {SESSION_TIMEOUT} Seconds - Please Try Again!")
|
|
except requests.exceptions.RequestException:
|
|
raise Exception("Error: Runware Request Failed!")
|
|
except Exception as e:
|
|
if "invalid api key" in str(e).lower():
|
|
PromptServer.instance.send_sync("runwareError", {
|
|
"success": False,
|
|
"errorMessage": str(e),
|
|
"errorCode": 401,
|
|
})
|
|
raise InterruptProcessingException()
|
|
else:
|
|
raise Exception(f"Error: {e}")
|
|
return False
|
|
|
|
def convertTensor2IMG(tensorImage):
|
|
imageNP = (tensorImage.squeeze().numpy() * 255).astype(np.uint8)
|
|
image = Image.fromarray(imageNP)
|
|
with io.BytesIO() as buffer:
|
|
image.save(buffer, format="webp")
|
|
imgb64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
|
|
imgDataURI = f'data:image/webp;base64,{imgb64}'
|
|
return imgDataURI
|
|
|
|
def convertIMG2Tensor(b64img):
|
|
imgbytes = base64.b64decode(b64img)
|
|
image = Image.open(io.BytesIO(imgbytes))
|
|
imageNP = np.array(image).astype(np.float32) / 255.0
|
|
tensorImage = torch.from_numpy(imageNP).squeeze()
|
|
return tensorImage
|
|
|
|
def convertImageB64List(imageDataObject):
|
|
images = ()
|
|
for result in imageDataObject["data"]:
|
|
generatedImage = result.get("imageBase64Data", result.get("maskImageBase64Data", result.get("guideImageBase64Data", False)))
|
|
generatedImage = convertIMG2Tensor(generatedImage)
|
|
images += (generatedImage,)
|
|
images = torch.stack(images, dim=0)
|
|
return (images,) |