Update app.py

app.py

@@ -9,17 +9,24 @@ from huggingface_hub import hf_hub_download
 import spaces
 from typing import Union, Sequence, Mapping, Any
 
-# Adicionar o caminho da pasta ComfyUI ao sys.path primeiro
+# 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)
 
-# Agora importar os módulos do ComfyUI
+# 2. Imports do ComfyUI
 import folder_paths
 from nodes import NODE_CLASS_MAPPINGS, init_extra_nodes
-import comfy.utils
 
-# Configuração inicial e diagnóstico CUDA
+# 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())
@@ -27,70 +34,68 @@ print("Quantidade de GPUs:", torch.cuda.device_count())
 if torch.cuda.is_available():
     print("GPU atual:", torch.cuda.get_device_name(0))
 
-# 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)
-
-# Inicializar nós extras
+# 5. Inicialização do ComfyUI
 print("Inicializando ComfyUI...")
 init_extra_nodes()
 
-# Helper function
+# 6. Helper Functions
 def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
     try:
         return obj[index]
     except KeyError:
         return obj["result"][index]
 
-# Baixar modelos
+# 7. Download de Modelos
 def download_models():
     print("Baixando modelos...")
     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", "models/diffusion_models"),
-        ("google/siglip-so400m-patch14-384", "model.safetensors", "models/clip_vision")
+        ("black-forest-labs/FLUX.1-Redux-dev", "flux1-redux-dev.safetensors", "style_models"),
+        ("comfyanonymous/flux_text_encoders", "t5xxl_fp16.safetensors", "text_encoders"),
+        ("zer0int/CLIP-GmP-ViT-L-14", "ViT-L-14-TEXT-detail-improved-hiT-GmP-HF.safetensors", "text_encoders"),
+        ("black-forest-labs/FLUX.1-dev", "ae.safetensors", "vae"),
+        ("black-forest-labs/FLUX.1-dev", "flux1-dev.safetensors", "diffusion_models"),
+        ("google/siglip-so400m-patch14-384", "model.safetensors", "clip_vision")
     ]
     
-    for repo_id, filename, local_dir in models:
+    for repo_id, filename, model_type in models:
         try:
-            os.makedirs(local_dir, exist_ok=True)
+            model_dir = os.path.join(models_dir, model_type)
+            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=local_dir)
+            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
 
-# Download models no início
+# 8. Download e Inicialização dos Modelos
+print("Baixando modelos...")
 download_models()
 
-# Inicializar modelos
 print("Inicializando modelos...")
 with torch.inference_mode():
     # CLIP
     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",
+        clip_name1="t5xxl_fp16.safetensors",
+        clip_name2="ViT-L-14-TEXT-detail-improved-hiT-GmP-HF.safetensors",
+        type="flux"
     )
 
     # Style Model
     stylemodelloader = NODE_CLASS_MAPPINGS["StyleModelLoader"]()
     stylemodelloader_441 = stylemodelloader.load_style_model(
-        style_model_name="models/style_models/flux1-redux-dev.safetensors"
+        style_model_name="flux1-redux-dev.safetensors"
     )
 
     # VAE
     vaeloader = NODE_CLASS_MAPPINGS["VAELoader"]()
     vaeloader_359 = vaeloader.load_vae(
-        vae_name="models/vae/ae.safetensors"
+        vae_name="ae.safetensors"
     )
 
+# 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)):
     try:
@@ -113,14 +118,6 @@ def generate_image(prompt, input_image, lora_weight, guidance, downsampling_fact
                 conditioning=encoded_text[0]
             )
 
-            # 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=stylemodelloader_441[0]
-            )
-
             # Redux Advanced
             reduxadvanced = NODE_CLASS_MAPPINGS["ReduxAdvanced"]()
             redux_result = reduxadvanced.apply_stylemodel(
@@ -150,7 +147,7 @@ def generate_image(prompt, input_image, lora_weight, guidance, downsampling_fact
                 sampler_name="euler",
                 scheduler="simple",
                 denoise=1,
-                model=lora_model[0],
+                model=stylemodelloader_441[0],
                 positive=redux_result[0],
                 negative=flux_guidance[0],
                 latent_image=empty_latent[0]
@@ -173,7 +170,7 @@ def generate_image(prompt, input_image, lora_weight, guidance, downsampling_fact
         print(f"Erro ao gerar imagem: {str(e)}")
         return None
 
-# Interface Gradio
+# 10. Interface Gradio
 with gr.Blocks() as app:
     gr.Markdown("# FLUX Redux Image Generator")
     
@@ -249,7 +246,7 @@ with gr.Blocks() as app:
             generate_btn = gr.Button("Generate Image")
         
         with gr.Column():
-            output_image = gr.Image(label="Generated Image", type="pil")
+            output_image = gr.Image(label="Generated Image", type="filepath")
     
     generate_btn.click(
         fn=generate_image,