Files

190 lines
8.2 KiB
Python

import os
import os.path
import sys
import json
import requests
import threading
import time
#import torch
from tqdm import tqdm
from huggingface_hub import snapshot_download
#import ctypes
import comfy
class FooocusInpaintWrapper:
def __init__(self):
self.node_dir = os.path.normpath(os.path.dirname(os.path.realpath(__file__)))
self.comfyui_dir = os.path.normpath(self.node_dir.split("custom_nodes")[0])
self.fooocus_dir = os.path.normpath(self.node_dir + "/Fooocus")
if not os.path.isfile(os.path.normpath(self.fooocus_dir + "/config.txt")):
config = {
"path_checkpoints": [
os.path.normpath(self.comfyui_dir + "/models/checkpoints")
],
"path_loras": [
os.path.normpath(self.comfyui_dir + "/models/loras")
],
"path_embeddings": os.path.normpath(self.comfyui_dir + "/models/embeddings"),
"path_vae_approx": os.path.normpath(self.comfyui_dir + "/models/vae_approx"),
"path_vae": os.path.normpath(self.comfyui_dir + "/models/vae"),
"path_upscale_models": os.path.normpath(self.comfyui_dir + "/models/upscale_/models"),
"path_inpaint": os.path.normpath(self.comfyui_dir + "/models/inpaint"),
"path_controlnet": os.path.normpath(self.comfyui_dir + "/models/controlnet"),
"path_clip_vision": os.path.normpath(self.comfyui_dir + "/models/clip_vision"),
"path_fooocus_expansion": os.path.normpath(self.comfyui_dir + "/models/prompt_expansion/fooocus_expansion"),
"path_wildcards": os.path.normpath(self.comfyui_dir + "/wildcards"),
"path_safety_checker": os.path.normpath(self.comfyui_dir + "/models/safety_checker"),
"path_sam": os.path.normpath(self.comfyui_dir + "/models/sam"),
"path_outputs": os.path.normpath(self.fooocus_dir + "/outputs")
}
with open(os.path.normpath(self.fooocus_dir + "/config.txt"), "w") as f:
json.dump(config, f, indent=4)
self.find_replace(os.path.normpath(self.fooocus_dir + "/modules/config.py"), './config.txt', os.path.normpath(self.fooocus_dir + "/config.txt"))
self.find_replace(os.path.normpath(self.fooocus_dir + "/modules/config.py"), './presets/default.json', os.path.normpath(self.fooocus_dir + "/presets/default.json"))
self.find_replace(os.path.normpath(self.fooocus_dir + "/args_manager.py"), 'args_parser.args = args_parser.parser.parse_args()', 'args_parser.args, unknown = args_parser.parser.parse_known_args()')
snapshot_download(repo_id="LykosAI/GPT-Prompt-Expansion-Fooocus-v2", local_dir = os.path.normpath(self.comfyui_dir + '/models/prompt_expansion/fooocus_expansion'))
self.file_download('https://huggingface.co/lllyasviel/misc/resolve/main/xlvaeapp.pth?download=true', os.path.normpath(self.comfyui_dir + '/models/vae_approx/xlvaeapp.pth'))
self.file_download('https://huggingface.co/lllyasviel/fav_models/resolve/main/fav/juggernautXL_v8Rundiffusion.safetensors?download=true', os.path.normpath(self.comfyui_dir + '/models/checkpoints/juggernautXL_v8Rundiffusion.safetensors'))
self.file_download('https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_offset_example-lora_1.0.safetensors?download=true', os.path.normpath(self.comfyui_dir + '/models/loras/sd_xl_offset_example-lora_1.0.safetensors'))
def file_download(self, url, save_path):
if os.path.isfile(save_path):
return True
response = requests.get(url, stream=True)
if response.status_code == 200:
total_size = int(response.headers.get('content-length', 0))
# Usa tqdm per mostrare la barra di avanzamento
with open(save_path, 'wb') as file, tqdm(
desc=save_path.split("models")[-1],
total=total_size,
unit='B',
unit_scale=True,
) as bar:
for chunk in response.iter_content(chunk_size=1024):
file.write(chunk)
bar.update(len(chunk))
else:
print(f"Error: {response.status_code}")
def find_replace(self, file_path, find_text, replace_text):
with open(file_path, 'r') as file:
content = file.read()
new_content = content.replace(find_text, replace_text.replace("\\","\\\\"))
with open(file_path, 'w') as file:
file.write(new_content)
@classmethod
def INPUT_TYPES(s):
def get_files(folder):
checkpoint_dir = os.path.normpath(os.path.dirname(os.path.realpath(__file__)).split("custom_nodes")[0] + "/models/" + folder)
if os.path.exists(checkpoint_dir):
files = [
f for f in os.listdir(checkpoint_dir)
if f.endswith(".ckpt") or f.endswith(".safetensors")
]
else:
files = []
return files if files else ["No checkpoints found"]
checkpoints = get_files("checkpoints")
loras = get_files("loras")
loras.append("None")
return {
"required": {
"image": ("IMAGE",),
"mask": ("MASK",),
"performance": (["Quality", "Speed", "Extreme Speed", "Lightning", "Hyper-SD"], {"default": "Speed"}),
"checkpoint": (checkpoints, {"default": checkpoints[0]}),
"prompt": ("STRING", {"multiline": True}),
"guidance_scale": ("FLOAT", {"default": 4, "min": 0, "max": 30.0}),
"image_sharpness": ("FLOAT", {"default": 2, "min": 0, "max": 30.0}),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "tooltip": "The random seed used for creating the noise."}),
},
"optional": {
"negative_prompt": ("STRING", {"multiline": True}),
"method": (["", "Improve Detail", "Modify content"],),
"inpaint_additional_prompt": ("STRING", {"multiline": True}),
"outpainting": (["", "Left", "Right", "Top", "Bottom", "Top + Bottom", "Left + Right"],),
"lora1": (loras, {"default": "None"}),
"lora1_weight": ("FLOAT", {"default": 0.1, "min": -2, "max": 2}),
"lora2": (loras, {"default": "None"}),
"lora2_weight": ("FLOAT", {"default": 0.1, "min": -2, "max": 2}),
"lora3": (loras, {"default": "None"}),
"lora3_weight": ("FLOAT", {"default": 0.1, "min": -2, "max": 2}),
"lora4": (loras, {"default": "None"}),
"lora4_weight": ("FLOAT", {"default": 0.1, "min": -2, "max": 2}),
"lora5": (loras, {"default": "None"}),
"lora5_weight": ("FLOAT", {"default": 0.1, "min": -2, "max": 2}),
},
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("image",)
FUNCTION = "start_inpaint"
#OUTPUT_NODE = False
CATEGORY = "inpaint"
def start_periodic_check(self, percentage_file):
pbar = comfy.utils.ProgressBar(100)
# Simuliamo un controllo periodico (ad esempio, ogni 1 secondo)
while True:
if os.path.exists(percentage_file):
with open(percentage_file, 'r') as f:
percentage = f.read().strip()
pbar.update_absolute(int(percentage))
else:
print(f"File {percentage_file} non trovato.")
time.sleep(5)
def start_inpaint(self, image, mask, performance, checkpoint, prompt, guidance_scale, image_sharpness, seed, negative_prompt, method, inpaint_additional_prompt, outpainting, lora1, lora1_weight, lora2, lora2_weight, lora3, lora3_weight, lora4, lora4_weight, lora5, lora5_weight):
def check_percentage():
self.start_periodic_check(os.path.normpath(self.node_dir+"/percentage.txt"))
try:
with open(os.path.normpath(self.node_dir + "/percentage.txt"), "w") as f:
f.write("0")
os.remove(os.path.normpath(self.node_dir + "/js/image.png"))
Image.new("RGB", (100, 100), (0, 0, 0)).save(os.path.normpath(self.node_dir + "/js/image.png"))
except:
pass
periodic_thread1 = threading.Thread(target=check_percentage)
periodic_thread1.start()
sys.path.append(self.fooocus_dir)
from launch import fooocusinpaintlaunch
new_image = fooocusinpaintlaunch(self.fooocus_dir, image, mask, performance, checkpoint, prompt, negative_prompt, guidance_scale, image_sharpness, seed, method, inpaint_additional_prompt, outpainting, lora1, lora1_weight, lora2, lora2_weight, lora3, lora3_weight, lora4, lora4_weight, lora5, lora5_weight)
os.remove(os.path.normpath(self.node_dir + "/js/image.png"))
return (new_image)
# A dictionary that contains all nodes you want to export with their names
# NOTE: names should be globally unique
NODE_CLASS_MAPPINGS = {
"FooocusInpaintWrapper": FooocusInpaintWrapper
}
# A dictionary that contains the friendly/humanly readable titles for the nodes
NODE_DISPLAY_NAME_MAPPINGS = {
"FirstNode": "Fooocus Inpaint Wrapper"
}