Update app.py

app.py

@@ -1,53 +1,25 @@
 import os
-import sys
 import random
+import sys
+from typing import Sequence, Mapping, Any, Union
 import torch
-from pathlib import Path
-from PIL import Image
 import gradio as gr
 from huggingface_hub import hf_hub_download
 import spaces
-from typing import Union, Sequence, Mapping, Any
-
-# 1. Configuração de Caminhos e Imports
-current_dir = os.path.dirname(os.path.abspath(__file__))
-comfyui_path = os.path.join(current_dir, "ComfyUI")
-sys.path.append(comfyui_path)
-
-# 2. Imports do ComfyUI
-import folder_paths
-from nodes import NODE_CLASS_MAPPINGS, init_extra_nodes
-
-# 3. Configuração de Diretórios
-BASE_DIR = os.path.dirname(os.path.realpath(__file__))
-output_dir = os.path.join(BASE_DIR, "output")
-models_dir = os.path.join(BASE_DIR, "models")
-os.makedirs(output_dir, exist_ok=True)
-os.makedirs(models_dir, exist_ok=True)
-folder_paths.set_output_directory(output_dir)
-
-# 4. Diagnóstico CUDA
-print("Python version:", sys.version)
-print("Torch version:", torch.__version__)
-print("CUDA disponível:", torch.cuda.is_available())
-print("Quantidade de GPUs:", torch.cuda.device_count())
-if torch.cuda.is_available():
-    print("GPU atual:", torch.cuda.get_device_name(0))
-
-# 5. Inicialização do ComfyUI
-print("Inicializando ComfyUI...")
-init_extra_nodes()
+from comfy import model_management
+from nodes import NODE_CLASS_MAPPINGS
+from PIL import Image
+from folder_paths import folder_paths
 
-# 6. Helper Functions
+# Helper function to get values from objects
 def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
     try:
         return obj[index]
     except KeyError:
         return obj["result"][index]
 
-# 7. Download de Modelos
+# Download models from Hugging Face
 def download_models():
-    print("Baixando modelos...")
     models = [
         ("black-forest-labs/FLUX.1-Redux-dev", "flux1-redux-dev.safetensors", "style_models"),
         ("comfyanonymous/flux_text_encoders", "t5xxl_fp16.safetensors", "text_encoders"),
@@ -63,46 +35,55 @@ def download_models():
             os.makedirs(model_dir, exist_ok=True)
             print(f"Baixando {filename} de {repo_id}...")
             hf_hub_download(repo_id=repo_id, filename=filename, local_dir=model_dir)
-            # Adicionar o diretório ao folder_paths
             folder_paths.add_model_folder_path(model_type, model_dir)
         except Exception as e:
             print(f"Erro ao baixar {filename} de {repo_id}: {str(e)}")
             continue
 
-# 8. Download e Inicialização dos Modelos
-download_models()
-
-print("Inicializando modelos...")
-with torch.inference_mode():
-    # CLIP
-    dualcliploader = NODE_CLASS_MAPPINGS["DualCLIPLoader"]()
-    dualcliploader_357 = dualcliploader.load_clip(
-        clip_name1="t5xxl_fp16.safetensors",
-        clip_name2="ViT-L-14-TEXT-detail-improved-hiT-GmP-HF.safetensors",
-        type="flux"
-    )
-
-    # CLIP Vision
-    clipvisionloader = NODE_CLASS_MAPPINGS["CLIPVisionLoader"]()
-    clip_vision = clipvisionloader.load_clip(
-        clip_name="model.safetensors"
-    )
-
-    # Style Model
-    stylemodelloader = NODE_CLASS_MAPPINGS["StyleModelLoader"]()
-    stylemodelloader_441 = stylemodelloader.load_style_model(
-        style_model_name="flux1-redux-dev.safetensors"
-    )
+# Configuração de Diretórios
+BASE_DIR = os.path.dirname(os.path.realpath(__file__))
+output_dir = os.path.join(BASE_DIR, "output")
+models_dir = os.path.join(BASE_DIR, "models")
+os.makedirs(output_dir, exist_ok=True)
+os.makedirs(models_dir, exist_ok=True)
+folder_paths.set_output_directory(output_dir)
 
-    # VAE
-    vaeloader = NODE_CLASS_MAPPINGS["VAELoader"]()
-    vaeloader_359 = vaeloader.load_vae(
-        vae_name="ae.safetensors"
-    )
+# Download and load models
+download_models()
 
-# 9. Função de Geração
-@spaces.GPU
-def generate_image(prompt, input_image, lora_weight, guidance, downsampling_factor, weight, seed, width, height, batch_size, steps, progress=gr.Progress(track_tqdm=True)):
+# Load models globally
+intconstant = NODE_CLASS_MAPPINGS["INTConstant"]()
+dualcliploader = NODE_CLASS_MAPPINGS["DualCLIPLoader"]()
+dualcliploader_357 = dualcliploader.load_clip(
+    clip_name1="t5xxl_fp16.safetensors",
+    clip_name2="ViT-L-14-TEXT-detail-improved-hiT-GmP-HF.safetensors",
+    type="flux"
+)
+clipvisionloader = NODE_CLASS_MAPPINGS["CLIPVisionLoader"]()
+clip_vision = clipvisionloader.load_clip(
+    clip_name="model.safetensors"
+)
+stylemodelloader = NODE_CLASS_MAPPINGS["StyleModelLoader"]()
+stylemodelloader_441 = stylemodelloader.load_style_model(
+    style_model_name="flux1-redux-dev.safetensors"
+)
+vaeloader = NODE_CLASS_MAPPINGS["VAELoader"]()
+vaeloader_359 = vaeloader.load_vae(
+    vae_name="ae.safetensors"
+)
+
+# Pre-load models
+model_loaders = [dualcliploader_357, vaeloader_359, clip_vision, 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)
+
+# Function to generate images
+@spaces.GPU(duration=60)  # Adjust duration as needed
+def generate_image(prompt, input_image, lora_weight, guidance, downsampling_factor, weight, seed, width, height, batch_size, steps):
     try:
         with torch.inference_mode():
             # Codificar texto
@@ -176,7 +157,7 @@ def generate_image(prompt, input_image, lora_weight, guidance, downsampling_fact
         print(f"Erro ao gerar imagem: {str(e)}")
         return None
 
-# 10. Interface Gradio
+# Gradio Interface
 with gr.Blocks() as app:
     gr.Markdown("# FLUX Redux Image Generator")