smthemex-ComfyUI_InstantIR_Wrapper/InstantIR_node.py

# !/usr/bin/env python
# -*- coding: UTF-8 -*-
import os
import torch
import numpy as np
import logging
import folder_paths

from .utils import load_list_images, tensor2imglist, tensor_upscale,instantIR_main, instantIR_load_model, auto_downlaod,clear_memory

node_cur_path = os.path.dirname(os.path.abspath(__file__))
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
if torch.cuda.is_available():
    torch_dtype = torch.float16
else:
    torch_dtype = torch.float32
    
MAX_SEED = np.iinfo(np.int32).max

# add checkpoints dir
InstantIR_current_path = os.path.join(folder_paths.models_dir, "InstantIR")
if not os.path.exists(InstantIR_current_path):
    os.makedirs(InstantIR_current_path)

try:
    folder_paths.add_model_folder_path("InstantIR", InstantIR_current_path, False)
except:
    folder_paths.add_model_folder_path("InstantIR", InstantIR_current_path)

InstantIR_base_path = os.path.join(InstantIR_current_path, "models")  # InstantIR/models
if not os.path.exists(InstantIR_base_path):
    os.makedirs(InstantIR_base_path)

InstantIR_dino_path = os.path.join(InstantIR_current_path, "dino")  # InstantIR/models
if not os.path.exists(InstantIR_dino_path):
    os.makedirs(InstantIR_dino_path)


class InstantIR_Loader:
    def __init__(self):
        pass
    
    @classmethod
    def INPUT_TYPES(s):
        INSTANT_LIST=["none"]+[i for i in folder_paths.get_filename_list("InstantIR") if ".bin" in i  or ".pt" in i] if folder_paths.get_filename_list("InstantIR") else ["none"]
        return {
            "required": {
                "sdxl_checkpoints": (["none"] + folder_paths.get_filename_list("checkpoints"),),
                "dino_repo": ("STRING", {"default": "facebook/dinov2-large"}),
                "adapter_checkpoints": (INSTANT_LIST,),
                "aggregator_checkpoints": (INSTANT_LIST,),
                "lora": (["none"] + folder_paths.get_filename_list("loras"),),
                "InstantIR_lora": (INSTANT_LIST,),
                "use_clip_encoder": ("BOOLEAN", {"default": False},),
            }
        }
    
    RETURN_TYPES = ("InstantIR_MODEL",)
    RETURN_NAMES = ("model",)
    FUNCTION = "main_"
    CATEGORY = "InstantIR"
    
    def main_(self, sdxl_checkpoints, dino_repo, adapter_checkpoints, aggregator_checkpoints,lora, InstantIR_lora,
             use_clip_encoder,):
        if not dino_repo:
            logging.info("no dino files in dir ,auto download from facebook/dinov2-large")
            vision_encoder_path = auto_downlaod(InstantIR_current_path, "dino")
        else:
            vision_encoder_path = dino_repo
        
        if sdxl_checkpoints != "none":
            sdxl_path = folder_paths.get_full_path("checkpoints", sdxl_checkpoints)
        else:
            #sdxl_path='stabilityai/stable-diffusion-xl-base-1.0'
            raise "need chocie a sdxl checkpoint"
        
        if adapter_checkpoints != "none":
            adapter_path = folder_paths.get_full_path("InstantIR", adapter_checkpoints)
        else:
            raise "need chocie a adapter checkpoint"
        
        if InstantIR_lora != "none":
            previewer_lora_path = folder_paths.get_full_path("InstantIR", InstantIR_lora)
        else:
            raise "need chocie a lora checkpoint"
        
        if lora != "none":
            lora_path = folder_paths.get_full_path("loras", lora)
        else:
            raise "need chocie a SDXL lcm lora checkpoint"
        
        if aggregator_checkpoints != "none":
            aggregator_path = folder_paths.get_full_path("InstantIR", aggregator_checkpoints)
        else:
            raise "need chocie a aggregator checkpoint"
        
        model = instantIR_load_model(use_clip_encoder, vision_encoder_path, sdxl_path, adapter_path,
                                     previewer_lora_path,lora_path, aggregator_path,device)
        
        logging.info("loading checkpoint done.")
        model.to(dtype=torch_dtype)
        if not torch.backends.mps.is_available():
            model.enable_model_cpu_offload()  # MPS not support
        return (model,)


class InstantIR_Sampler:
    def __init__(self):
        pass
    
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("InstantIR_MODEL",),
                "pixels": ("IMAGE",),
                "prompt": ("STRING", {"multiline": True, "default": "Photorealistic, highly detailed, hyper detailed photo - realistic maximum detail, 32k, ultra HD, extreme meticulous detailing, skin pore detailing, hyper sharpness, perfect without deformations, taken using a Canon EOS R camera, Cinematic, High Contrast, Color Grading. "}),
                "negative_prompt": ("STRING", {"multiline": True,
                                               "default": "blurry, out of focus, unclear, depth of field, over-smooth, sketch, oil painting, cartoon, CG Style, 3D render, unreal engine, dirty, messy, worst quality, low quality, frames, painting, illustration, drawing, art, watermark, signature, jpeg artifacts, deformed, lowres"}),
                "seed": ("INT", {"default": 0, "min": 0, "max": MAX_SEED}),
                "steps": ("INT", {"default": 20, "min": 1, "max": 100}),
                "cfg": ("FLOAT", {"default": 7, "min": 0.1, "max": 10.0, "step": 0.1, "round": 0.01}),
                "creative_restoration": ("BOOLEAN", {"default": False},),
                "width": ("INT", {"default": 768, "min": 128, "max": 1024, "step": 64, "display": "number"}),
                "height": ("INT", {"default": 768, "min": 128, "max": 1024, "step": 64, "display": "number"}),
                "preview_start": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 30.0, "step": 0.1,}),
                "guidance_end": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 30.0, "step": 0.1,}),
                "batch_size": ("INT", {"default": 1, "min": 1, "max": 64, "step": 1, }),
            },
        }
    
    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "main"
    CATEGORY = "InstantIR"
    
    def main(self, model, pixels, prompt, negative_prompt, seed, steps, cfg, creative_restoration, width, height, preview_start,guidance_end,batch_size):
        # pre image
        image = tensor_upscale(pixels, width, height)
        image_list, _ = tensor2imglist(image, np_out=True)
        
        logging.info(f"Start infer {len(image_list)} images.")
        ouput_img = instantIR_main(image_list, model, seed, creative_restoration, steps, prompt, negative_prompt, cfg,
                                   batch_size, device,preview_start,guidance_end)
        logging.info("finish processing")
        clear_memory()
        image = load_list_images(ouput_img)
        
        return (image,)


NODE_CLASS_MAPPINGS = {
    "InstantIR_Loader": InstantIR_Loader,
    "InstantIR_Sampler": InstantIR_Sampler,
}

# A dictionary that contains the friendly/humanly readable titles for the nodes
NODE_DISPLAY_NAME_MAPPINGS = {
    "InstantIR_Loader": "InstantIR_Loader",
    "InstantIR_Sampler": "InstantIR_Sampler",
}