Update app.py

app.py

@@ -1,25 +1,29 @@
 import os
+import sys
 import random
 import torch
 from pathlib import Path
 from PIL import Image
 import gradio as gr
 from huggingface_hub import hf_hub_download
-import sys
+import spaces
+from typing import Union, Sequence, Mapping, Any
 
-# 1. Configuração de Caminhos
+# 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
-from nodes import NODE_CLASS_MAPPINGS
 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
@@ -29,102 +33,90 @@ 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("GPU não disponível. Usando CPU.")
 
-# 5. Download de Modelos
+# 5. Inicialização do ComfyUI
+print("Inicializando ComfyUI...")
+init_extra_nodes()
+
+# 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]
+
+# 7. Download de Modelos
 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"),
-        ("zer0int/CLIP-GmP-ViT-L-14", "ViT-L-14-TEXT-detail-improved-hiT-GmP-TE-only-HF.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.sft", "diffusion_models"),
-        ("google/siglip-so400m-patch14-384", "model.safetensors", "clip_vision"),
-        ("black-forest-labs/FLUX.1-Redux-dev", "NFTNIK_FLUX.1[dev]_LoRA.safetensors", "lora")
+        ("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:
-        model_dir = os.path.join(BASE_DIR, "models", 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)
+        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)
+            # 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
 
-# 6. Load custom nodes
-def import_custom_nodes():
-    import asyncio
-    import execution
-    from nodes import init_extra_nodes
-    import server
+# 8. Download e Inicialização dos Modelos
+print("Baixando modelos...")
+download_models()
 
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
+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"
+    )
 
-    server_instance = server.PromptServer(loop)
-    execution.PromptQueue(server_instance)
-    init_extra_nodes()
+    # CLIP Vision
+    clipvisionloader = NODE_CLASS_MAPPINGS["CLIPVisionLoader"]()
+    clip_vision = clipvisionloader.load_clip(
+        clip_name="model.safetensors"
+    )
 
-# 7. Main function to execute the workflow and generate an image
-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():
-            device = "cuda" if torch.cuda.is_available() else "cpu"
-            print(f"Using device: {device}")
+    # Style Model
+    stylemodelloader = NODE_CLASS_MAPPINGS["StyleModelLoader"]()
+    stylemodelloader_441 = stylemodelloader.load_style_model(
+        style_model_name="flux1-redux-dev.safetensors"
+    )
 
-            # Load CLIP
-            dualcliploader = NODE_CLASS_MAPPINGS["DualCLIPLoader"]()
-            dualcliploader_loaded = dualcliploader.load_clip(
-                clip_name1="t5xxl_fp16.safetensors",
-                clip_name2="ViT-L-14-TEXT-detail-improved-hiT-GmP-TE-only-HF.safetensors",
-                type="flux",
-                device=device
-            )
+    # VAE
+    vaeloader = NODE_CLASS_MAPPINGS["VAELoader"]()
+    vaeloader_359 = vaeloader.load_vae(
+        vae_name="ae.safetensors"
+    )
 
-            # Text Encoding
+# 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:
+        with torch.inference_mode():
+            # Codificar texto
             cliptextencode = NODE_CLASS_MAPPINGS["CLIPTextEncode"]()
             encoded_text = cliptextencode.encode(
                 text=prompt,
-                clip=dualcliploader_loaded[0]
+                clip=dualcliploader_357[0]
             )
 
-            # Load Style Model
-            stylemodelloader = NODE_CLASS_MAPPINGS["StyleModelLoader"]()
-            style_model = stylemodelloader.load_style_model(
-                style_model_name="flux1-redux-dev.safetensors"
-            )
-
-            # Load CLIP Vision
-            clipvisionloader = NODE_CLASS_MAPPINGS["CLIPVisionLoader"]()
-            clip_vision = clipvisionloader.load_clip(
-                clip_name="model.safetensors"
-            )
-
-            # Load Input Image
+            # Carregar e processar imagem
             loadimage = NODE_CLASS_MAPPINGS["LoadImage"]()
             loaded_image = loadimage.load_image(image=input_image)
 
-            # Load VAE
-            vaeloader = NODE_CLASS_MAPPINGS["VAELoader"]()
-            vae = vaeloader.load_vae(vae_name="ae.safetensors")
-
-            # Load UNET
-            unetloader = NODE_CLASS_MAPPINGS["UNETLoader"]()
-            unet = unetloader.load_unet(
-                unet_name="flux1-dev.sft",
-                weight_dtype="fp8_e4m3fn"
-            )
-
-            # 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[0]
-            )
-
             # Flux Guidance
             fluxguidance = NODE_CLASS_MAPPINGS["FluxGuidance"]()
             flux_guidance = fluxguidance.append(
@@ -139,14 +131,13 @@ 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=style_model[0],
+                style_model=stylemodelloader_441[0],
                 clip_vision=clip_vision[0],
                 image=loaded_image[0]
             )
 
-            # Empty Latent Image
+            # Empty Latent
             emptylatentimage = NODE_CLASS_MAPPINGS["EmptyLatentImage"]()
             empty_latent = emptylatentimage.generate(
                 width=width,
@@ -163,49 +154,30 @@ 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]
             )
 
-            # VAE Decode
+            # Decodificar VAE
             vaedecode = NODE_CLASS_MAPPINGS["VAEDecode"]()
             decoded = vaedecode.decode(
                 samples=sampled[0],
-                vae=vae[0]
-            )
-
-            # Save the image in the output directory
-            saveimage = NODE_CLASS_MAPPINGS["SaveImage"]()
-            temp_filename = f"Flux_{random.randint(0, 99999)}"
-            saveimage.save_images(
-                filename_prefix=temp_filename,
-                images=decoded[0]
+                vae=vaeloader_359[0]
             )
 
-            # Add a delay to ensure the file system updates
-            import time
-            time.sleep(0.5)
-
-            # Dynamically retrieve the correct file name
-            saved_files = [f for f in os.listdir(output_dir) if f.startswith(temp_filename)]
-            if not saved_files:
-                raise FileNotFoundError(f"Output file not found: Expected files starting with {temp_filename}")
-
-            # Get the full path of the saved file
-            temp_path = os.path.join(output_dir, saved_files[0])
-            print(f"Image saved at: {temp_path}")
-
-            # Return the saved image for Gradio display
-            output_image = Image.open(temp_path)
-            return output_image
+            # 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)
 
+            return temp_path
     except Exception as e:
-        print(f"Error during generation: {str(e)}")
+        print(f"Erro ao gerar imagem: {str(e)}")
         return None
 
-# 8. Gradio Interface
+# 10. Interface Gradio
 with gr.Blocks() as app:
     gr.Markdown("# FLUX Redux Image Generator")
     
@@ -281,7 +253,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,
@@ -302,5 +274,4 @@ with gr.Blocks() as app:
     )
 
 if __name__ == "__main__":
-    # Download_models()
     app.launch(share=True)