diff --git a/topaz_upscaler.py b/topaz_upscaler.py index 9c42905..937ae60 100644 --- a/topaz_upscaler.py +++ b/topaz_upscaler.py @@ -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'] \ No newline at end of file +NODE_CLASS_MAPPINGS = {"TopazUpscaler": TopazUpscaler} +NODE_DISPLAY_NAME_MAPPINGS = {"TopazUpscaler": "Topaz Upscaler"} +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']