[
  {
    "path": ".github/workflows/publish.yml",
    "content": "name: Publish to Comfy registry\non:\n  workflow_dispatch:\n  push:\n    branches:\n      - main\n      - master\n    paths:\n      - \"pyproject.toml\"\n\njobs:\n  publish-node:\n    name: Publish Custom Node to registry\n    runs-on: ubuntu-latest\n    steps:\n      - name: Check out code\n        uses: actions/checkout@v4\n      - name: Publish Custom Node\n        uses: Comfy-Org/publish-node-action@main\n        with:\n          ## Add your own personal access token to your Github Repository secrets and reference it here.\n          personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}\n"
  },
  {
    "path": ".gitignore",
    "content": "__pycache__"
  },
  {
    "path": "LICENSE",
    "content": "MIT License\n\nCopyright (c) 2024 choey\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE."
  },
  {
    "path": "README.md",
    "content": "# What is Comfy-Topaz?\nComfy-Topaz is a custom node for ComfyUI, which integrates with Topaz Photo AI to enhance (upscale, sharpen, denoise, etc.) images, allowing this traditionally asynchronous step to become a part of ComfyUI workflows.\n\n# Requirements\n- Licensed installation of [Topaz Photo AI](https://www.topazlabs.com/downloads): This provides `tpai.exe`, the path to which should be set for the `Topaz Photo AI (tpai.exe)` setting in ComfyUI.\n\n# Installation\nClone this repo into `ComfyUI/custom_nodes` and restart ComfyUI.\n\n# Usage\n## Auto-Pilot Settings\nThis is the simplest use case, which relies on Topaz Photo AI to auto-detect and apply those settings. This is done by omitting the `upscale` and `sharpen` settings as inputs:\n\n![simple demo showing the auto-pilot settings](demo1.png)\n\nSometimes auto-pilot settings don't yield the best results, which warrants manual tuning.\n\nOn the output side, `autopilot_settings` shows what the auto-pilot settings were, and `settings` shows all the features used and knobs turned to generate the final image.\n\n## Manual Settings\nOverride auto-pilot settings by providing manual settings:\n\n![demo showing the upscale settings override](demo2.png)\n\nA good starting point is by copying over params from auto-pilot from which to iterate. I copied over the `denoise`, `deblur`, and `detail` values and changed the model from `Standard V2` to `High Fidelity`.\n\n# TODO\n- Output `*Settings` nodes rather than json, to eliminate the manual steps of copying values when overriding settings.\n- Add a button to run auto-pilot analysis without applying the settings on the image.\n- Map `param1`, `param2`, `param3`, ... to the actual param name (e.g., `denoise`, `deblur`, and `detail` for Upscale Settings)\n- Expose more settings (denoise, face recovery, text recovery, WB/exposure adjustment)."
  },
  {
    "path": "__init__.py",
    "content": "from .topaz import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS\nimport os\nimport shutil\nimport __main__\n\nWEB_DIRECTORY = \"./web\"\n__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY']\n\n# ensure extensions_path exists\nextentions_path = os.path.join(os.path.dirname(os.path.realpath(__main__.__file__)), \"web\", \"extensions\", \"topaz\")\nif not os.path.exists(extentions_path):\n    os.makedirs(extentions_path)\n\n# copy all *.js files from js_path to extesnions_path\njs_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), \"web\", \"js\")\nfor file in os.listdir(js_path):\n    if file.endswith(\".js\"):\n        src_file = os.path.join(js_path, file)\n        dst_file = os.path.join(extentions_path, file)\n        if os.path.exists(dst_file):\n            os.remove(dst_file)\n        shutil.copy(src_file, dst_file)\n        print('installed %s to %s' % (file, extentions_path))"
  },
  {
    "path": "install.sh",
    "content": "#!/bin/sh\n# copy js to extensions directory (useful for development)\nmkdir -p ../../web/extensions/topaz\ncp -r ./web/js/* ../../web/extensions/topaz/ "
  },
  {
    "path": "pyproject.toml",
    "content": "[project]\nname = \"comfy-topaz\"\ndescription = \"Comfy-Topaz is a custom node for ComfyUI, which integrates with Topaz Photo AI to enhance (upscale, sharpen, denoise, etc.) images, allowing this traditionally asynchronous step to become a part of ComfyUI workflows.\\nNOTE: Requires licensed installation of Topaz Photo AI\"\nversion = \"1.0.1\"\nlicense = { file = \"LICENSE\" }\n\n[project.urls]\nRepository = \"https://github.com/choey/Comfy-Topaz\"\n#  Used by Comfy Registry https://comfyregistry.org\n\n[tool.comfy]\nPublisherId = \"choey\"\nDisplayName = \"Comfy-Topaz\"\nIcon = \"\"\n"
  },
  {
    "path": "topaz.py",
    "content": "import numpy as np\nimport os\nimport pprint\nimport time\nimport folder_paths\nimport torch\nimport subprocess\nimport json\n\nfrom PIL import Image, ImageOps\nfrom typing import Optional\nimport json\n\nclass TopazUpscaleSettings:\n    @classmethod\n    def INPUT_TYPES(cls):\n        return {\n            'required': {\n                'enabled': (['true', 'false'], {'default': 'true'}),\n                'model': ([\n                    'Standard', \n                    'Standard V2', \n                    'High Fidelity', \n                    'High Fidelity V2', \n                    #'Graphics', what is this\n                    'Low Resolution'\n                ], {'default': 'Standard V2'}),\n                'scale': ('FLOAT', {'default': 2.0, 'min': 0, 'max': 10, 'round': False, }),\n                'denoise': ('FLOAT', {'default': 0.2, 'min': 0, 'max': 10, 'round': False, 'display': 'denoise (param1)'}),\n                'deblur': ('FLOAT', {'default': 0.2, 'min': 0, 'max': 10, 'round': False, 'display': 'deblur (param2)'}),\n                'detail': ('FLOAT', {'default': 0.2, 'min': 0, 'max': 10, 'round': False, 'display': 'detail (param3)'}),\n            },\n            'optional': {\n\n            },\n        }\n\n    RETURN_TYPES = ('TopazUpscaleSettings',)\n    RETURN_NAMES = ('upscale_settings',)\n    FUNCTION = 'init'\n    CATEGORY = 'image'\n    OUTPUT_NODE = False\n    OUTPUT_IS_LIST = (False,)\n    \n    def init(self, enabled, model, scale, denoise, deblur, detail):\n        self.enabled = str(True).lower() == enabled.lower()\n        self.model = model\n        self.scale = scale\n        self.denoise = denoise\n        self.deblur = deblur\n        self.detail = detail\n        return (self,)\n\nclass TopazSharpenSettings:       \n    @classmethod\n    def INPUT_TYPES(cls):\n        return {\n            'required': {\n                'enabled': (['true', 'false'], {'default': 'true'}),\n                'model': ([\n                    'Standard', \n                    'Strong', \n                    # TODO: why don't these work?\n                    #'Lens Blur', \n                    #'Motion Blur', \n                ], {'default': 'Standard'}),\n                'compression': ('FLOAT', {'default': 0.5, 'min': 0, 'max': 1, 'round': 0.01,}),\n                'is_lens': (['true', 'false'], {'default': 'false'}),\n                'lensblur': ('FLOAT', {'default': 0.0, 'min': 0, 'max': 10, 'round': False,}),\n                'mask': (['true', 'false'], {'default': 'false'}),\n                'motionblur': ('FLOAT', {'default': 0.0, 'min': 0, 'max': 10, 'round': False,}),\n                'noise': ('FLOAT', {'default': 0.0, 'min': 0, 'max': 10, 'round': False,}),\n                'strength': ('FLOAT', {'default': 0.0, 'min': 0, 'max': 10, 'round': False, 'display': 'strength (param1)'}), # TODO: why doesn't \"display\" work?\n                'denoise': ('FLOAT', {'default': 0.0, 'min': 0, 'max': 10, 'round': False, \"display\": 'denoise (param2)'}),   # param2 (Lens/Motion Blur only)\n            },\n            'optional': {\n                \n            },\n        }\n\n    RETURN_TYPES = ('TopazSharpenSettings',)\n    RETURN_NAMES = ('sharpen_settings',)\n    FUNCTION = 'init'\n    CATEGORY = 'image'\n    OUTPUT_IS_LIST = (False,)\n    \n    def init(self, enabled, model, compression, is_lens, lensblur, mask, motionblur, noise, strength, denoise):\n        self.enabled = str(True).lower() == enabled.lower()\n        self.model = model\n        self.compression = compression\n        self.is_lens = is_lens\n        self.lensblur = lensblur\n        self.mask = mask\n        self.motionblur = motionblur\n        self.noise = noise\n        self.strength = strength\n        self.denoise = denoise\n        return (self,)\n\nclass TopazPhotoAI:\n    '''\n    A node that uses Topaz Image AI (tpai.exe) behind the scenes to enhance (upscale/sharpen/denoise/etc.) the given image(s).\n    \n    If no settings are provided, auto-detected (auto-pilot) settings are used.\n    '''\n    def __init__(self):\n        self.this_dir = os.path.dirname(os.path.abspath(__file__))\n        self.comfy_dir = os.path.abspath(os.path.join(self.this_dir, '..', '..'))\n        self.subfolder = 'upscaled'\n        self.output_dir = os.path.join(self.comfy_dir, 'temp')\n        self.prefix = 'tpai'\n        # self.tpai = 'C:/Program Files/Topaz Labs LLC/Topaz Photo AI/tpai.exe'\n\n    @classmethod\n    def INPUT_TYPES(cls):\n        return {\n            'required': {\n                'images': ('IMAGE',),\n            },\n            'optional': {\n                'compression': ('INT', {\n                    'default': 2,\n                    'min': 0,\n                    'max': 10,\n                }),\n                'tpai_exe': ('STRING', {\n                    'default': '',                    \n                }),\n                # 'blur_level': ('FLOAT', {'default': -1, 'min': -10, 'max': 10}),\n                # 'noise_level': ('FLOAT', {'default': -1, 'min': -10, 'max': 10}),\n                'upscale': ('TopazUpscaleSettings',),\n                'sharpen': ('TopazSharpenSettings',),\n            },\n            \"hidden\": {\n            }\n        }\n\n    RETURN_TYPES = ('STRING', 'STRING', 'IMAGE')\n    RETURN_NAMES = ('settings', 'autopilot_settings', 'IMAGE')\n    FUNCTION = 'upscale_image'\n    CATEGORY = 'image'\n    OUTPUT_NODE = True\n    OUTPUT_IS_LIST = (True, True, True)\n\n    def save_image(self, img, output_dir, filename):\n        if not os.path.exists(output_dir):\n            os.makedirs(output_dir)\n        file_path = os.path.join(output_dir, filename)\n        img.save(file_path)\n        return file_path\n\n    def load_image(self, image):\n        image_path = folder_paths.get_annotated_filepath(image)\n        i = Image.open(image_path)\n        i = ImageOps.exif_transpose(i)\n        image = i.convert('RGB')\n        image = np.array(image).astype(np.float32) / 255.0\n        image = torch.from_numpy(image)[None,]\n        return image\n\n    def get_settings(self, stdout):\n        '''\n        Extracts the settings JSON string from the stdout of the tpai.exe process\n        '''        \n        # find index of 'Final Settings for' in stdout\n        settings_start = stdout.find('Final Settings for')\n        # starting from settings_start, find the opening curly brace '{'\n        settings_start = stdout.find('{', settings_start)\n        # for each character after the opening curly brace, count the opening and closing curly braces\n        # when the count is zero, that is the end of the JSON string\n        # (escaped, mismatched braces shouldn't be a problem)\n        count = 0\n        settings_end = settings_start\n        for i in range(settings_start, len(stdout)):            \n            if stdout[i] == '{':\n                count += 1\n            elif stdout[i] == '}':\n                count -= 1\n            if count == 0:\n                settings_end = i\n                break\n        settings_json = str(stdout[settings_start : settings_end + 1])\n        settings = json.loads(settings_json)\n        autopilot_settings = settings.pop('autoPilotSettings')\n        user_settings_json = json.dumps(settings, indent=2).replace('\"', \"'\")\n        autopilot_settings_json = json.dumps(autopilot_settings, indent=2).replace('\"', \"'\")\n        \n        return user_settings_json, autopilot_settings_json\n\n    def topaz_upscale(self, img_file, compression=0, format='png', tpai_exe=None, \n                      upscale: Optional[TopazUpscaleSettings]=None, \n                      sharpen: Optional[TopazSharpenSettings]=None):\n        if not os.path.exists(tpai_exe):\n            raise ValueError('Topaz AI Upscaler not found at %s' % tpai_exe)\n        if compression < 0 or compression > 10:\n            raise ValueError('compression must be between 0 and 10')        \n        \n        target_dir = os.path.join(self.output_dir, self.subfolder)\n        tpai_args = [\n            tpai_exe,\n            '--output',        # output directory\n            target_dir,\n            '--compression',   # compression=[0,10] (default=2)\n            str(compression),\n            '--format',        # output format (omit to preserve original)\n            format,\n            '--showSettings',  # Prints out the final settings used when processing.\n        ]\n        \n        if upscale:\n            print('\\033[31mComfy-Topaz:\\033[0m upscaler override:', pprint.pformat(upscale))\n            tpai_args.append('--upscale')\n            if upscale.enabled:\n                tpai_args.append('%s=%g' % ('scale', upscale.scale))\n                tpai_args.append('%s=%g' % ('param1', upscale.denoise)) # Minor Denoise\n                tpai_args.append('%s=%g' % ('param2', upscale.deblur))  # Minor Deblur\n                tpai_args.append('%s=%g' % ('param3', upscale.detail))  # Fix Compression\n                tpai_args.append('%s=%s' % ('model', upscale.model))\n            else:\n                tpai_args.append('enabled=false')\n                \n            \n        if sharpen:\n            print('\\033[31mComfy-Topaz:\\033[0m sharpen override:', pprint.pformat(sharpen))\n            tpai_args.append('--sharpen')\n            if sharpen.enabled:\n                tpai_args.append('%s=Sharpen %s' % ('model', sharpen.model))\n                tpai_args.append('%s=%g' % ('compression', sharpen.compression))\n                tpai_args.append('%s=%s' % ('is_lens', sharpen.is_lens))\n                tpai_args.append('%s=%g' % ('lensblur', sharpen.lensblur))\n                tpai_args.append('%s=%s' % ('mask', sharpen.mask))\n                tpai_args.append('%s=%g' % ('motionblur', sharpen.motionblur))\n                tpai_args.append('%s=%g' % ('noise', sharpen.noise))\n                tpai_args.append('%s=%g' % ('param1', sharpen.strength))\n                tpai_args.append('%s=%g' % ('param2', sharpen.denoise))\n            else:\n                tpai_args.append('enabled=false')\n            \n        tpai_args.append(img_file)\n        print('\\033[31mComfy-Topaz:\\033[0m tpaie.exe args:', pprint.pformat(tpai_args))\n        p_tpai = subprocess.run(tpai_args, capture_output=True, text=True, shell=False)\n        print('\\033[31mComfy-Topaz:\\033[0m tpaie.exe return code:', p_tpai.returncode)\n        print('\\033[31mComfy-Topaz:\\033[0m tpaie.exe STDOUT:', p_tpai.stdout)\n        print('\\033[31mComfy-Topaz:\\033[0m tpaie.exe STDERR:', p_tpai.stderr)\n\n        user_settings, autopilot_settings = self.get_settings(p_tpai.stdout)\n\n        return (os.path.join(target_dir, os.path.basename(img_file)), user_settings, autopilot_settings)\n\n    def upscale_image(self, images, compression=0, format='png', tpai_exe=None, \n                      upscale: Optional[TopazUpscaleSettings]=None, \n                      sharpen: Optional[TopazSharpenSettings]=None):\n        now_millis = int(time.time() * 1000)\n        prefix = '%s-%d' % (self.prefix, now_millis)\n        upscaled_images = []\n        upscale_user_settings = []\n        upscale_autopilot_settings = []\n        count = 0\n        for image in images:\n            count += 1\n            i = 255.0 * image.cpu().numpy()\n            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))\n            img_file = self.save_image(\n                img, self.output_dir, '%s-%d.png' % (prefix, count)\n            )\n            (upscaled_img_file, user_settings, autopilot_settings) = self.topaz_upscale(img_file, compression, format, tpai_exe=tpai_exe, upscale=upscale, sharpen=sharpen)\n            upscaled_image = self.load_image(upscaled_img_file)\n            upscaled_images.append(upscaled_image)\n            upscale_user_settings.append(user_settings)\n            upscale_autopilot_settings.append(autopilot_settings)\n\n        return (upscale_user_settings, upscale_autopilot_settings, upscaled_images)\n\nNODE_CLASS_MAPPINGS = {\n    'TopazPhotoAI': TopazPhotoAI,\n    'TopazSharpenSettings': TopazSharpenSettings,\n    'TopazUpscaleSettings': TopazUpscaleSettings,\n}\n\nNODE_DISPLAY_NAME_MAPPINGS = {\n    'TopazPhotoAI': 'Topaz Photo AI',\n    'TopazSharpenSettings': 'Topaz Sharpen Settings',\n    'TopazUpscaleSettings': 'Topaz Upscale Settings',\n}\n"
  },
  {
    "path": "web/js/topaz.js",
    "content": "import { app } from \"../../scripts/app.js\";\nimport { ComfyWidgets } from \"../../scripts/widgets.js\";\n\nlet tpai_setting;\nconst id = \"comfy.topaz\";\nconst ext = {\n    name: id,\n    async setup(app) {\n        tpai_setting = app.ui.settings.addSetting({\n            id,\n            name: \"Topaz Photo AI (tpai.exe)\",\n            defaultValue: \"C:\\\\Program Files\\\\Topaz Labs LLC\\\\Topaz Photo AI\\\\tpai.exe\",\n            type: \"string\",\n        });        \n    },\n    async beforeRegisterNodeDef(nodeType, nodeData, _app) {\n        if (nodeData.name === 'TopazPhotoAI') {\n            const ensureTpai = async (node) => {\n                const tpaiWidget = node.widgets.find(w => w.name === \"tpai_exe\");\n                if (tpaiWidget && tpaiWidget.value === \"\") {\n                    tpaiWidget.value = tpai_setting.value;\n                }\n            }\n\n            const onConfigure = nodeType.prototype.onConfigure;\n            nodeType.prototype.onConfigure = function () {\n                const r = onConfigure ? onConfigure.apply(this, arguments) : undefined;\n                ensureTpai(this);\n                return r;\n            };\n\n            const onNodeCreated = nodeType.prototype.onNodeCreated;\n            nodeType.prototype.onNodeCreated = function () {\n                const r = onNodeCreated ? onNodeCreated.apply(this, arguments) : undefined;\n                ensureTpai(this);\n                return r;\n            };\n        }\n    },    \n}\napp.registerExtension(ext);"
  }
]