Update app.py

app.py

@@ -6,6 +6,7 @@ from PIL import Image
 import gradio as gr
 from huggingface_hub import hf_hub_download
 from nodes import NODE_CLASS_MAPPINGS
+from comfy import model_management
 import folder_paths
 
 # Diretório base e de saída
@@ -14,39 +15,61 @@ output_dir = os.path.join(BASE_DIR, "output")
 os.makedirs(output_dir, exist_ok=True)
 folder_paths.set_output_directory(output_dir)
 
-# Baixar os modelos necessários
-
+# Baixar e carregar os modelos necessários
 hf_hub_download(repo_id="black-forest-labs/FLUX.1-Redux-dev",
-filename="flux1-redux-dev.safetensors",
-local_dir="models/style_models")
+                filename="flux1-redux-dev.safetensors",
+                local_dir="models/style_models")
 
 hf_hub_download(repo_id="comfyanonymous/flux_text_encoders",
-filename="t5xxl_fp16.safetensors",
-local_dir="models/text_encoders")
+                filename="t5xxl_fp16.safetensors",
+                local_dir="models/text_encoders")
 
 hf_hub_download(repo_id="zer0int/CLIP-GmP-ViT-L-14",
-filename="ViT-L-14-TEXT-detail-improved-hiT-GmP-HF.safetensors",
-local_dir="models/text_encoders")
+                filename="ViT-L-14-TEXT-detail-improved-hiT-GmP-HF.safetensors",
+                local_dir="models/text_encoders")
 
 hf_hub_download(repo_id="black-forest-labs/FLUX.1-dev",
-filename="ae.safetensors",
-local_dir="models/vae")
+                filename="ae.safetensors",
+                local_dir="models/vae")
 
 hf_hub_download(repo_id="black-forest-labs/FLUX.1-dev",
-filename="flux1-dev.safetensors.safetensors",
-local_dir="models/diffusion_models")
+                filename="flux1-dev.safetensors.safetensors",
+                local_dir="models/diffusion_models")
 
 hf_hub_download(repo_id="google/siglip-so400m-patch14-384",
-filename="model.safetensors",
-local_dir="models/clip_vision")
+                filename="model.safetensors",
+                local_dir="models/clip_vision")
 
 hf_hub_download(repo_id="nftnik/NFTNIK-FLUX.1-dev-LoRA",
-filename="NFTNIK_FLUX.1[dev]_LoRA.safetensors",
-local_dir="models/lora")
+                filename="NFTNIK_FLUX.1[dev]_LoRA.safetensors",
+                local_dir="models/lora")
+
+# Inicializar os nós e pré-carregar os modelos
+intconstant = NODE_CLASS_MAPPINGS["INTConstant"]()
+dualcliploader = NODE_CLASS_MAPPINGS["DualCLIPLoader"]()
+dualcliploader_357 = dualcliploader.load_clip(
+    clip_name1="models/text_encoders/t5xxl_fp16.safetensors",
+    clip_name2="models/text_encoders/ViT-L-14-TEXT-detail-improved-hiT-GmP-HF.safetensors",
+    type="flux",
+)
+stylemodelloader = NODE_CLASS_MAPPINGS["StyleModelLoader"]()
+stylemodelloader_441 = stylemodelloader.load_style_model(
+    style_model_name="models/style_models/flux1-redux-dev.safetensors"
+)
+vaeloader = NODE_CLASS_MAPPINGS["VAELoader"]()
+vaeloader_359 = vaeloader.load_vae(vae_name="models/vae/ae.safetensors")
+
+# Lista de modelos para carregamento na GPU
+model_loaders = [dualcliploader_357, vaeloader_359, stylemodelloader_441]
+valid_models = [
+    getattr(loader[0], 'patcher', loader[0]) 
+    for loader in model_loaders
+    if not isinstance(loader[0], dict) and not isinstance(getattr(loader[0], 'patcher', None), dict)
+]
+model_management.load_models_gpu(valid_models)
 
 # Função para importar nodes personalizados
 def import_custom_nodes():
-    """Carregar nodes customizados."""
     import asyncio
     import execution
     from nodes import init_extra_nodes
@@ -62,49 +85,32 @@ def import_custom_nodes():
 # Função principal de geração
 def generate_image(prompt, input_image, lora_weight, guidance, downsampling_factor, weight, seed, width, height, batch_size, steps):
     import_custom_nodes()
-
     try:
         with torch.inference_mode():
-            # Carregar CLIP
-            dualcliploader = NODE_CLASS_MAPPINGS["DualCLIPLoader"]()
-            dualcliploader_loaded = dualcliploader.load_clip(
-                clip_name1="models/text_encoders/t5xxl_fp16.safetensors",
-                clip_name2="models/clip_vision/ViT-L-14-TEXT-detail-improved-hiT-GmP-HF.safetensors",
-                type="flux"
-            )
-
             # Codificar texto
             cliptextencode = NODE_CLASS_MAPPINGS["CLIPTextEncode"]()
             encoded_text = cliptextencode.encode(
                 text=prompt,
-                clip=dualcliploader_loaded[0]
+                clip=dualcliploader_357[0]
             )
 
-            # Carregar modelos de estilo e LoRA
-            stylemodelloader = NODE_CLASS_MAPPINGS["StyleModelLoader"]()
-            style_model = stylemodelloader.load_style_model(
-                style_model_name="models/style_models/flux1-redux-dev.safetensors"
-            )
+            # Carregar LoRA
             loraloadermodelonly = NODE_CLASS_MAPPINGS["LoraLoaderModelOnly"]()
             lora_model = loraloadermodelonly.load_lora_model_only(
                 lora_name="models/lora/NFTNIK_FLUX.1[dev]_LoRA.safetensors",
                 strength_model=lora_weight,
-                model=style_model[0]
+                model=stylemodelloader_441[0]
             )
 
             # Processar imagem de entrada
             loadimage = NODE_CLASS_MAPPINGS["LoadImage"]()
             loaded_image = loadimage.load_image(image=input_image)
 
-            # Configurações adicionais e saída
-            vaeloader = NODE_CLASS_MAPPINGS["VAELoader"]()
-            vae = vaeloader.load_vae(vae_name="models/vae/ae.safetensors")
-
             # Decodificar e salvar
             vaedecode = NODE_CLASS_MAPPINGS["VAEDecode"]()
             decoded = vaedecode.decode(
                 samples=lora_model[0],
-                vae=vae[0]
+                vae=vaeloader_359[0]
             )
 
             temp_filename = f"Flux_{random.randint(0, 99999)}.png"
@@ -124,14 +130,6 @@ with gr.Blocks() as app:
             prompt_input = gr.Textbox(label="Prompt", placeholder="Digite seu prompt aqui...", lines=5)
             input_image = gr.Image(label="Imagem de Entrada", type="filepath")
             lora_weight = gr.Slider(minimum=0, maximum=2, step=0.1, value=0.6, label="Peso LoRA")
-            guidance = gr.Slider(minimum=0, maximum=20, step=0.1, value=3.5, label="Orientação")
-            downsampling_factor = gr.Slider(minimum=1, maximum=8, step=1, value=3, label="Fator de Redução")
-            weight = gr.Slider(minimum=0, maximum=2, step=0.1, value=1.0, label="Peso do Modelo")
-            seed = gr.Number(value=random.randint(1, 2**64), label="Seed", precision=0)
-            width = gr.Number(value=1024, label="Largura", precision=0)
-            height = gr.Number(value=1024, label="Altura", precision=0)
-            batch_size = gr.Number(value=1, label="Tamanho do Lote", precision=0)
-            steps = gr.Number(value=20, label="Etapas", precision=0)
             generate_btn = gr.Button("Gerar Imagem")
 
         with gr.Column():
@@ -139,19 +137,7 @@ with gr.Blocks() as app:
 
     generate_btn.click(
         fn=generate_image,
-        inputs=[
-            prompt_input,
-            input_image,
-            lora_weight,
-            guidance,
-            downsampling_factor,
-            weight,
-            seed,
-            width,
-            height,
-            batch_size,
-            steps
-        ],
+        inputs=[prompt_input, input_image, lora_weight],
         outputs=[output_image]
     )