Update topaz_upscaler.py
This commit is contained in:
+186
-65
@@ -8,6 +8,56 @@ import io
|
||||
|
||||
BASE_URL = "https://api.topazlabs.com/image/v1"
|
||||
|
||||
# All available modes/endpoints
|
||||
TOPAZ_MODES = [
|
||||
"enhance",
|
||||
"sharpen",
|
||||
"denoise",
|
||||
"restore",
|
||||
"lighting"
|
||||
]
|
||||
|
||||
# All available models (full list from API docs)
|
||||
TOPAZ_MODELS = [
|
||||
# Enhance - Standard
|
||||
"Standard V2",
|
||||
"Low Resolution V2",
|
||||
"CGI",
|
||||
"High Fidelity V2",
|
||||
"Text Refine",
|
||||
# Enhance - Generative
|
||||
"Redefine",
|
||||
"Recovery V2",
|
||||
"Standard MAX",
|
||||
"Wonder",
|
||||
# Sharpen - Standard
|
||||
"Standard", # Sharpen
|
||||
"Strong", # Sharpen
|
||||
"Lens Blur",
|
||||
"Lens Blur V2",
|
||||
"Motion Blur",
|
||||
"Natural", # Sharpen
|
||||
"Refocus",
|
||||
# Sharpen - Generative
|
||||
"Super Focus V2",
|
||||
# Denoise
|
||||
"Normal",
|
||||
"Strong", # Denoise
|
||||
"Extreme",
|
||||
# Restore
|
||||
"Dust-Scratch",
|
||||
# Lighting
|
||||
"Adjust",
|
||||
"White Balance"
|
||||
]
|
||||
|
||||
# Format mappings for download accept header
|
||||
FORMAT_ACCEPT = {
|
||||
"jpeg": "image/jpeg",
|
||||
"png": "image/png",
|
||||
"tiff": "image/tiff"
|
||||
}
|
||||
|
||||
class TopazUpscaler:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
@@ -19,105 +69,176 @@ class TopazUpscaler:
|
||||
"default": "",
|
||||
"placeholder": "Enter your Topaz API key"
|
||||
}),
|
||||
"mode": (TOPAZ_MODES, {"default": "enhance"}),
|
||||
"model": (TOPAZ_MODELS, {"default": "Standard V2"}),
|
||||
},
|
||||
"optional": {
|
||||
"output_width": ("INT", {
|
||||
"default": 3840,
|
||||
"min": 512,
|
||||
"max": 16000,
|
||||
"step": 64
|
||||
"default": 0, # 0 means auto
|
||||
"min": 1,
|
||||
"max": 32000,
|
||||
"step": 64,
|
||||
"display": "number"
|
||||
}),
|
||||
"output_height": ("INT", {
|
||||
"default": 0, # 0 means auto
|
||||
"min": 1,
|
||||
"max": 32000,
|
||||
"step": 64,
|
||||
"display": "number"
|
||||
}),
|
||||
"crop_to_fill": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
"optional": {}
|
||||
"output_format": (["jpeg", "png", "tiff"], {"default": "jpeg"}),
|
||||
"face_enhancement": ("BOOLEAN", {"default": True}),
|
||||
"denoise_strength": ("FLOAT", {
|
||||
"default": 0.5,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.05
|
||||
}),
|
||||
"sharpen_strength": ("FLOAT", {
|
||||
"default": 0.5,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.05
|
||||
}),
|
||||
"strength": ("FLOAT", {
|
||||
"default": 0.5,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.05
|
||||
}),
|
||||
"fix_compression": ("FLOAT", {
|
||||
"default": 0.0,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.05
|
||||
}),
|
||||
"timeout_seconds": ("INT", {
|
||||
"default": 300,
|
||||
"min": 60,
|
||||
"max": 1800,
|
||||
"step": 30
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("upscaled_image",)
|
||||
FUNCTION = "upscale"
|
||||
CATEGORY = "image/upscaling"
|
||||
RETURN_NAMES = ("processed_image",)
|
||||
FUNCTION = "process"
|
||||
CATEGORY = "image/topaz"
|
||||
|
||||
def _submit_job(self, image_path, api_key, output_width, crop_to_fill):
|
||||
headers = {
|
||||
"X-API-Key": api_key,
|
||||
"accept": "application/json"
|
||||
}
|
||||
def _get_endpoint_path(self, mode):
|
||||
if mode == "restore":
|
||||
return "/restore-gen/async"
|
||||
elif mode == "enhance":
|
||||
return "/enhance/async"
|
||||
elif mode == "sharpen":
|
||||
return "/sharpen/async"
|
||||
elif mode == "denoise":
|
||||
return "/denoise/async"
|
||||
elif mode == "lighting":
|
||||
return "/lighting/async"
|
||||
else:
|
||||
raise ValueError(f"Unsupported mode: {mode}")
|
||||
|
||||
def _submit_job(self, image_path, api_key, mode, params):
|
||||
path = self._get_endpoint_path(mode)
|
||||
headers = {"X-API-Key": api_key, "accept": "application/json"}
|
||||
with open(image_path, 'rb') as f:
|
||||
files = {'image': f}
|
||||
data = {
|
||||
"output_width": output_width,
|
||||
"crop_to_fill": str(crop_to_fill).lower(),
|
||||
"output_format": "jpeg"
|
||||
}
|
||||
response = requests.post(f"{BASE_URL}/enhanceAsync", headers=headers, files=files, data=data)
|
||||
response = requests.post(f"{BASE_URL}{path}", headers=headers, files=files, data=params)
|
||||
response.raise_for_status()
|
||||
return response.json()["task_id"]
|
||||
data = response.json()
|
||||
return data["process_id"], data.get("eta", 0)
|
||||
|
||||
def _wait_for_completion(self, task_id, api_key, timeout=300):
|
||||
def _wait_for_completion(self, process_id, api_key, timeout):
|
||||
headers = {"X-API-Key": api_key}
|
||||
start_time = time.time()
|
||||
while time.time() - start_time < timeout:
|
||||
resp = requests.get(f"{BASE_URL}/getStatus?task_id={task_id}", headers=headers)
|
||||
start = time.time()
|
||||
while time.time() - start < timeout:
|
||||
resp = requests.get(f"{BASE_URL}/getStatus?process_id={process_id}", headers=headers)
|
||||
resp.raise_for_status()
|
||||
status = resp.json()
|
||||
if status["status"] == "completed":
|
||||
return True
|
||||
elif status["status"] == "failed":
|
||||
raise Exception(f"Topaz task {task_id} failed: {status.get('error', 'Unknown error')}")
|
||||
if status["status"] == "failed":
|
||||
raise Exception(f"Topaz task failed: {status.get('error', 'Unknown')}")
|
||||
time.sleep(5)
|
||||
raise TimeoutError(f"Topaz task {task_id} timed out after {timeout} seconds.")
|
||||
raise TimeoutError(f"Task {process_id} timed out after {timeout}s")
|
||||
|
||||
def _download_result(self, task_id, api_key):
|
||||
headers = {
|
||||
"X-API-Key": api_key,
|
||||
"accept": "image/jpeg"
|
||||
}
|
||||
resp = requests.get(f"{BASE_URL}/getDownloadOutput?task_id={task_id}", headers=headers)
|
||||
def _download_result(self, process_id, api_key, output_format):
|
||||
accept = FORMAT_ACCEPT.get(output_format, "image/jpeg")
|
||||
headers = {"X-API-Key": api_key, "accept": accept}
|
||||
resp = requests.get(f"{BASE_URL}/getDownloadOutput?process_id={process_id}", headers=headers)
|
||||
resp.raise_for_status()
|
||||
return resp.content
|
||||
|
||||
def upscale(self, image, api_key, output_width, crop_to_fill):
|
||||
if not api_key or api_key.strip() == "":
|
||||
raise ValueError("Topaz API key is required. Please enter it in the node.")
|
||||
def process(self, image, api_key, mode, model,
|
||||
output_width=0, output_height=0, crop_to_fill=False,
|
||||
output_format="jpeg", face_enhancement=True,
|
||||
denoise_strength=0.5, sharpen_strength=0.5, strength=0.5,
|
||||
fix_compression=0.0, timeout_seconds=300):
|
||||
|
||||
# Convert ComfyUI image tensor [B,H,W,C] float32 [0,1] -> PIL Image
|
||||
if not api_key.strip():
|
||||
raise ValueError("Topaz API key is required.")
|
||||
|
||||
# Build params dict
|
||||
params = {
|
||||
"model": model,
|
||||
"output_format": output_format,
|
||||
"face_enhancement": str(face_enhancement).lower(),
|
||||
"denoise_strength": denoise_strength,
|
||||
"sharpen_strength": sharpen_strength,
|
||||
"strength": strength,
|
||||
"fix_compression": fix_compression,
|
||||
}
|
||||
|
||||
# Enhance-specific params
|
||||
if mode == "enhance":
|
||||
if output_width > 0:
|
||||
params["output_width"] = output_width
|
||||
if output_height > 0:
|
||||
params["output_height"] = output_height
|
||||
params["crop_to_fill"] = str(crop_to_fill).lower()
|
||||
|
||||
# Convert ComfyUI image to temp file
|
||||
img_np = (image[0].cpu().numpy() * 255).astype(np.uint8)
|
||||
pil_image = Image.fromarray(img_np)
|
||||
|
||||
# Save to temp file
|
||||
with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp_input:
|
||||
pil_image.save(tmp_input.name, format="JPEG", quality=95)
|
||||
input_path = tmp_input.name
|
||||
pil_img = Image.fromarray(img_np)
|
||||
suffix = f".{output_format}" if output_format != "jpeg" else ".jpg"
|
||||
with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
|
||||
if output_format == "tiff":
|
||||
pil_img.save(tmp.name, format="TIFF")
|
||||
elif output_format == "png":
|
||||
pil_img.save(tmp.name, format="PNG")
|
||||
else:
|
||||
pil_img.save(tmp.name, format="JPEG", quality=95)
|
||||
input_path = tmp.name
|
||||
|
||||
try:
|
||||
print(f"[Topaz Upscaler] Submitting job with width={output_width}, crop={crop_to_fill}")
|
||||
task_id = self._submit_job(input_path, api_key, output_width, crop_to_fill)
|
||||
print(f"[Topaz Upscaler] Task ID: {task_id}")
|
||||
print(f"[Topaz] Submitting {mode} job with model={model}")
|
||||
process_id, eta = self._submit_job(input_path, api_key, mode, params)
|
||||
print(f"[Topaz] Process ID: {process_id}, ETA: {eta}")
|
||||
|
||||
print(f"[Topaz Upscaler] Waiting for completion...")
|
||||
self._wait_for_completion(task_id, api_key)
|
||||
print(f"[Topaz] Waiting for completion...")
|
||||
self._wait_for_completion(process_id, api_key, timeout_seconds)
|
||||
|
||||
print(f"[Topaz Upscaler] Downloading result...")
|
||||
result_bytes = self._download_result(task_id, api_key)
|
||||
|
||||
# Load result
|
||||
print(f"[Topaz] Downloading result...")
|
||||
result_bytes = self._download_result(process_id, api_key, output_format)
|
||||
result_pil = Image.open(io.BytesIO(result_bytes))
|
||||
|
||||
# Convert to ComfyUI tensor (drop alpha if present)
|
||||
result_np = np.array(result_pil).astype(np.float32) / 255.0
|
||||
result_tensor = result_np[np.newaxis, ...] # Add batch dim
|
||||
if len(result_np.shape) == 3 and result_np.shape[2] == 4:
|
||||
result_np = result_np[:, :, :3] # RGB only
|
||||
result_tensor = result_np[np.newaxis, ...] # Batch dim
|
||||
|
||||
return (result_tensor,)
|
||||
|
||||
finally:
|
||||
# Clean up input temp file
|
||||
if os.path.exists(input_path):
|
||||
os.unlink(input_path)
|
||||
|
||||
|
||||
# Node registration
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"TopazUpscaler": TopazUpscaler
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"TopazUpscaler": "Topaz Upscaler"
|
||||
}
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
NODE_CLASS_MAPPINGS = {"TopazUpscaler": TopazUpscaler}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {"TopazUpscaler": "Topaz Upscaler"}
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
|
||||
Reference in New Issue
Block a user