Update app.py

app.py

@@ -1,109 +1,107 @@
 import os
-import random
 import sys
-from typing import Sequence, Mapping, Any, Union
+import random
 import torch
+from pathlib import Path
+from PIL import Image
 import gradio as gr
 from huggingface_hub import hf_hub_download
 import spaces
-from PIL import Image
-from folder_paths import folder_paths
+from typing import Union, Sequence, Mapping, Any
+import folder_paths
+from nodes import NODE_CLASS_MAPPINGS, init_extra_nodes
+from comfy import model_management
+
+# 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)
 
-# Append ComfyUI path to sys.path
-current_dir = os.path.dirname(os.path.abspath(__file__))
-comfyui_path = os.path.join(current_dir, "ComfyUI")
-sys.path.append(comfyui_path)
+# 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))
 
-# Import Comfy modules
-from comfy import model_management
-from nodes import NODE_CLASS_MAPPINGS
+# Inicializar nós extras
+print("Inicializando ComfyUI...")
+init_extra_nodes()
 
-# Helper function to get values from objects
+# Helper function
 def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
     try:
         return obj[index]
     except KeyError:
         return obj["result"][index]
 
-# Download models from Hugging Face
-def download_models():
-    models = [
-        ("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, model_type in models:
-        try:
-            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=model_dir)
-            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
-
-# 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)
+# Inicializar modelos
+print("Inicializando modelos...")
+with torch.inference_mode():
+    # CLIP
+    dualcliploader = NODE_CLASS_MAPPINGS["DualCLIPLoader"]()
+    CLIP_MODEL = dualcliploader.load_clip(
+        clip_name1="t5xxl_fp16.safetensors",
+        clip_name2="ViT-L-14-TEXT-detail-improved-hiT-GmP-TE-only-HF.safetensors",
+        type="flux"
+    )
+
+    # Style Model
+    stylemodelloader = NODE_CLASS_MAPPINGS["StyleModelLoader"]()
+    STYLE_MODEL = stylemodelloader.load_style_model(
+        style_model_name="flux1-redux-dev.safetensors"
+    )
+
+    # VAE
+    vaeloader = NODE_CLASS_MAPPINGS["VAELoader"]()
+    VAE_MODEL = vaeloader.load_vae(
+        vae_name="ae.safetensors"
+    )
 
-# Download and load models
-download_models()
-
-# 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]) 
+    # UNET
+    unetloader = NODE_CLASS_MAPPINGS["UNETLoader"]()
+    UNET_MODEL = unetloader.load_unet(
+        unet_name="flux1-dev.sft",
+        weight_dtype="fp8_e4m3fn"
+    )
+
+    # CLIP Vision
+    clipvisionloader = NODE_CLASS_MAPPINGS["CLIPVisionLoader"]()
+    CLIP_VISION = clipvisionloader.load_clip(
+        clip_name="sigclip_vision_patch14_384.safetensors"
+    )
+
+model_loaders = [CLIP_MODEL, VAE_MODEL, UNET_MODEL, CLIP_VISION]
+model_management.load_models_gpu([
+    loader[0].patcher if hasattr(loader[0], 'patcher') else 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):
+@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:
         with torch.inference_mode():
-            # Codificar texto
+            # Text Encoding
             cliptextencode = NODE_CLASS_MAPPINGS["CLIPTextEncode"]()
             encoded_text = cliptextencode.encode(
                 text=prompt,
-                clip=dualcliploader_357[0]
+                clip=CLIP_MODEL[0]
             )
 
-            # Carregar e processar imagem
+            # Load Input Image
             loadimage = NODE_CLASS_MAPPINGS["LoadImage"]()
             loaded_image = loadimage.load_image(image=input_image)
 
+            # Load LoRA
+            loraloadermodelonly = NODE_CLASS_MAPPINGS["LoraLoaderModelOnly"]()
+            lora_model = loraloadermodelonly.load_lora_model_only(
+                lora_name="NFTNIK_FLUX.1[dev]_LoRA.safetensors",
+                strength_model=lora_weight,
+                model=UNET_MODEL[0]
+            )
+
             # Flux Guidance
             fluxguidance = NODE_CLASS_MAPPINGS["FluxGuidance"]()
             flux_guidance = fluxguidance.append(
@@ -118,13 +116,14 @@ def generate_image(prompt, input_image, lora_weight, guidance, downsampling_fact
                 downsampling_function="area",
                 mode="keep aspect ratio",
                 weight=weight,
+                autocrop_margin=0.1,
                 conditioning=flux_guidance[0],
-                style_model=stylemodelloader_441[0],
-                clip_vision=clip_vision[0],
+                style_model=STYLE_MODEL[0],
+                clip_vision=CLIP_VISION[0],
                 image=loaded_image[0]
             )
 
-            # Empty Latent
+            # Empty Latent Image
             emptylatentimage = NODE_CLASS_MAPPINGS["EmptyLatentImage"]()
             empty_latent = emptylatentimage.generate(
                 width=width,
@@ -141,17 +140,17 @@ def generate_image(prompt, input_image, lora_weight, guidance, downsampling_fact
                 sampler_name="euler",
                 scheduler="simple",
                 denoise=1,
-                model=stylemodelloader_441[0],
+                model=lora_model[0],
                 positive=redux_result[0],
                 negative=flux_guidance[0],
                 latent_image=empty_latent[0]
             )
 
-            # Decodificar VAE
+            # VAE Decode
             vaedecode = NODE_CLASS_MAPPINGS["VAEDecode"]()
             decoded = vaedecode.decode(
                 samples=sampled[0],
-                vae=vaeloader_359[0]
+                vae=VAE_MODEL[0]
             )
 
             # Salvar imagem
@@ -160,11 +159,12 @@ def generate_image(prompt, input_image, lora_weight, guidance, downsampling_fact
             Image.fromarray((decoded[0] * 255).astype("uint8")).save(temp_path)
 
             return temp_path
+
     except Exception as e:
         print(f"Erro ao gerar imagem: {str(e)}")
         return None
 
-# Gradio Interface
+# Interface Gradio
 with gr.Blocks() as app:
     gr.Markdown("# FLUX Redux Image Generator")
     
@@ -244,7 +244,19 @@ 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,
+            guidance,
+            downsampling_factor,
+            weight,
+            seed,
+            width,
+            height,
+            batch_size,
+            steps
+        ],
         outputs=[output_image]
     )