241 lines
9.5 KiB
Python
241 lines
9.5 KiB
Python
import os
|
|
import os.path
|
|
import sys
|
|
import json
|
|
import requests
|
|
import threading
|
|
import time
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
from tqdm import tqdm
|
|
from huggingface_hub import snapshot_download
|
|
import cv2
|
|
import ctypes
|
|
|
|
import comfy
|
|
|
|
class ImageUpdater:
|
|
def __init__(self, image_file, refresh_rate=5):
|
|
self.image_file = image_file
|
|
self.refresh_rate = refresh_rate # Tempo in secondi tra aggiornamenti
|
|
self.running = False
|
|
self.thread = None
|
|
|
|
def start(self):
|
|
"""Avvia il thread dell'aggiornamento immagine."""
|
|
if self.running:
|
|
return # Evita di avviare più thread
|
|
self.running = True
|
|
self.thread = threading.Thread(target=self.start_loop, daemon=True)
|
|
self.thread.start()
|
|
|
|
def start_loop(self):
|
|
"""Loop di aggiornamento immagini eseguito nel thread."""
|
|
cv2.namedWindow("Preview Inpainting", cv2.WINDOW_NORMAL) # Finestra ridimensionabile
|
|
cv2.resizeWindow("Preview Inpainting", 240, 256) # Imposta dimensione finestra
|
|
|
|
while self.running:
|
|
if os.path.exists(self.image_file):
|
|
img = Image.open(self.image_file)
|
|
img = np.array(img)
|
|
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
|
|
|
|
cv2.imshow("Preview Inpainting", img)
|
|
key = cv2.waitKey(1)
|
|
|
|
time.sleep(self.refresh_rate)
|
|
|
|
try:
|
|
cv2.destroyAllWindows() # Chiude la finestra alla fine del thread
|
|
except:
|
|
pass
|
|
|
|
def stop(self):
|
|
"""Ferma il thread e chiude la finestra."""
|
|
self.running = False
|
|
if self.thread:
|
|
self.thread.join() # Aspetta la chiusura del thread
|
|
|
|
try:
|
|
cv2.destroyAllWindows()
|
|
except:
|
|
pass
|
|
|
|
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")
|
|
Image.new("RGB", (100, 100), (0, 0, 0)).save(os.path.normpath(self.node_dir + "/image.png"))
|
|
except:
|
|
pass
|
|
|
|
periodic_thread1 = threading.Thread(target=check_percentage)
|
|
periodic_thread1.start()
|
|
|
|
updater = ImageUpdater(self.node_dir+"/image.png", refresh_rate=5)
|
|
updater.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)
|
|
|
|
updater.stop()
|
|
|
|
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"
|
|
} |