Update app.py

app.py

@@ -1,86 +1,92 @@
 import os
 import sys
-
-# Adicionar o caminho da pasta ComfyUI ao sys.path
-current_dir = os.path.dirname(os.path.abspath(__file__))
-comfyui_path = os.path.join(current_dir, "ComfyUI")
-sys.path.append(comfyui_path)
-
 import random
 import torch
 from pathlib import Path
 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
+import spaces
+from typing import Union, Sequence, Mapping, Any
 
+# Configuração inicial e 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))
+else:
+    print("CUDA não está disponível. Verificando por que:")
+    try:
+        torch.cuda.init()
+    except Exception as e:
+        print("Erro ao inicializar CUDA:", str(e))
 
-# Diretório base e de saída
+# Adicionar o caminho da pasta ComfyUI ao sys.path
+current_dir = os.path.dirname(os.path.abspath(__file__))
+comfyui_path = os.path.join(current_dir, "ComfyUI")
+sys.path.append(comfyui_path)
+
+from nodes import NODE_CLASS_MAPPINGS
+from comfy import model_management
+import folder_paths
+
+# Configuração de diretórios
 BASE_DIR = os.path.dirname(os.path.realpath(__file__))
 output_dir = os.path.join(BASE_DIR, "output")
 os.makedirs(output_dir, exist_ok=True)
 folder_paths.set_output_directory(output_dir)
 
-# 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")
-
-hf_hub_download(repo_id="comfyanonymous/flux_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")
-
-hf_hub_download(repo_id="black-forest-labs/FLUX.1-dev",
-                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")
-
-hf_hub_download(repo_id="google/siglip-so400m-patch14-384",
-                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")
-
-# 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
+# Helper function
+def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
+    try:
+        return obj[index]
+    except KeyError:
+        return obj["result"][index]
+
+# Baixar modelos necessários
+def download_models():
+    models = [
+        ("black-forest-labs/FLUX.1-Redux-dev", "flux1-redux-dev.safetensors", "models/style_models"),
+        ("comfyanonymous/flux_text_encoders", "t5xxl_fp16.safetensors", "models/text_encoders"),
+        ("zer0int/CLIP-GmP-ViT-L-14", "ViT-L-14-TEXT-detail-improved-hiT-GmP-HF.safetensors", "models/text_encoders"),
+        ("black-forest-labs/FLUX.1-dev", "ae.safetensors", "models/vae"),
+        ("black-forest-labs/FLUX.1-dev", "flux1-dev.safetensors.safetensors", "models/diffusion_models"),
+        ("google/siglip-so400m-patch14-384", "model.safetensors", "models/clip_vision"),
+        ("nftnik/NFTNIK-FLUX.1-dev-LoRA", "NFTNIK_FLUX.1[dev]_LoRA.safetensors", "models/lora")
+    ]
+    
+    for repo_id, filename, local_dir in models:
+        hf_hub_download(repo_id=repo_id, filename=filename, local_dir=local_dir)
+
+# Inicializar modelos
+print("Inicializando modelos...")
+with torch.inference_mode():
+    # Initialize nodes
+    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")
+
+    # Carregar modelos 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)
+
 def import_custom_nodes():
     import asyncio
     import execution
@@ -89,13 +95,13 @@ def import_custom_nodes():
 
     loop = asyncio.new_event_loop()
     asyncio.set_event_loop(loop)
-
     server_instance = server.PromptServer(loop)
     execution.PromptQueue(server_instance)
     init_extra_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):
+@spaces.GPU
+def generate_image(prompt, input_image, lora_weight, progress=gr.Progress(track_tqdm=True)):
+    """Função principal de geração com monitoramento de progresso"""
     import_custom_nodes()
     try:
         with torch.inference_mode():
@@ -103,7 +109,7 @@ def generate_image(prompt, input_image, lora_weight, guidance, downsampling_fact
             cliptextencode = NODE_CLASS_MAPPINGS["CLIPTextEncode"]()
             encoded_text = cliptextencode.encode(
                 text=prompt,
-                clip=dualcliploader_357[0]
+                clip=get_value_at_index(dualcliploader_357, 0)
             )
 
             # Carregar LoRA
@@ -111,23 +117,24 @@ def generate_image(prompt, input_image, lora_weight, guidance, downsampling_fact
             lora_model = loraloadermodelonly.load_lora_model_only(
                 lora_name="models/lora/NFTNIK_FLUX.1[dev]_LoRA.safetensors",
                 strength_model=lora_weight,
-                model=stylemodelloader_441[0]
+                model=get_value_at_index(stylemodelloader_441, 0)
             )
 
-            # Processar imagem de entrada
+            # Processar imagem
             loadimage = NODE_CLASS_MAPPINGS["LoadImage"]()
             loaded_image = loadimage.load_image(image=input_image)
 
-            # Decodificar e salvar
+            # Decodificar
             vaedecode = NODE_CLASS_MAPPINGS["VAEDecode"]()
             decoded = vaedecode.decode(
-                samples=lora_model[0],
-                vae=vaeloader_359[0]
+                samples=get_value_at_index(lora_model, 0),
+                vae=get_value_at_index(vaeloader_359, 0)
             )
 
+            # Salvar imagem
             temp_filename = f"Flux_{random.randint(0, 99999)}.png"
             temp_path = os.path.join(output_dir, temp_filename)
-            Image.fromarray((decoded[0] * 255).astype("uint8")).save(temp_path)
+            Image.fromarray((get_value_at_index(decoded, 0) * 255).astype("uint8")).save(temp_path)
 
             return temp_path
     except Exception as e:
@@ -154,4 +161,7 @@ with gr.Blocks() as app:
     )
 
 if __name__ == "__main__":
-    app.launch()
+    # Download models at startup
+    download_models()
+    # Launch the app
+    app.launch()