Update app.py

app.py

@@ -3,13 +3,18 @@ import sys
 import random
 import torch
 from pathlib import Path
-from PIL import Image
+import numpy as np
 import gradio as gr
 from huggingface_hub import hf_hub_download
 import spaces
 from typing import Union, Sequence, Mapping, Any
 import logging
 
+# Adicione se ainda não tiver
+from nodes import NODE_CLASS_MAPPINGS, init_extra_nodes, SaveImage  # <-- Node SaveImage
+from comfy import model_management
+import folder_paths
+
 # Configurar logging para debug
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
@@ -18,11 +23,6 @@ logger = logging.getLogger(__name__)
 current_dir = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(current_dir)
 
-# 2. Imports do ComfyUI
-import folder_paths
-from nodes import NODE_CLASS_MAPPINGS, init_extra_nodes
-from comfy import model_management
-
 # 3. Configuração de Diretórios
 BASE_DIR = os.path.dirname(os.path.realpath(__file__))
 output_dir = os.path.join(BASE_DIR, "output")
@@ -109,8 +109,6 @@ except Exception as e:
 # 8. Inicialização dos Modelos
 logger.info("Inicializando modelos...")
 try:
-    # Use torch.no_grad() em vez de torch.inference_mode() 
-    # para evitar o erro de version counter.
     with torch.no_grad():
         # CLIP
         logger.info("Carregando CLIP...")
@@ -155,7 +153,7 @@ try:
         unetloader = NODE_CLASS_MAPPINGS["UNETLoader"]()
         UNET_MODEL = unetloader.load_unet(
             unet_name="flux1-dev.safetensors",
-            weight_dtype="fp8_e4m3fn"  # Ajuste a seu hardware, se necessário
+            weight_dtype="fp8_e4m3fn"  # ajuste se preciso
         )
         if UNET_MODEL is None:
             raise ValueError("Falha ao carregar UNET model")
@@ -175,14 +173,13 @@ except Exception as e:
 @spaces.GPU
 def generate_image(
     prompt, input_image, lora_weight, guidance, downsampling_factor,
-    weight, seed, width, height, batch_size, steps, 
+    weight, seed, width, height, batch_size, steps,
     progress=gr.Progress(track_tqdm=True)
 ):
     try:
-        # Aqui também: no_grad() para evitar cálculo de gradientes
         with torch.no_grad():
             logger.info(f"Iniciando geração com prompt: {prompt}")
-           
+
             # Codificar texto
             cliptextencode = NODE_CLASS_MAPPINGS["CLIPTextEncode"]()
             encoded_text = cliptextencode.encode(
@@ -217,7 +214,7 @@ def generate_image(
                 image=loaded_image[0]
             )
 
-            # Criar latente vazio
+            # Empty Latent
             emptylatentimage = NODE_CLASS_MAPPINGS["EmptyLatentImage"]()
             empty_latent = emptylatentimage.generate(
                 width=width,
@@ -249,16 +246,37 @@ def generate_image(
                 vae=VAE_MODEL[0]
             )
 
-            # Salvar imagem
-            temp_filename = f"Flux_{random.randint(0, 99999)}.png"
-            temp_path = os.path.join(output_dir, temp_filename)
-            try:
-                Image.fromarray((decoded[0] * 255).astype("uint8")).save(temp_path)
-                logger.info(f"Imagem salva em: {temp_path}")
-                return temp_path
-            except Exception as e:
-                logger.error(f"Erro ao salvar imagem: {str(e)}")
-                return None
+            # ======================== SALVAR IMAGEM USANDO O NODE SaveImage ======================
+            logger.info("Salvando imagem via node SaveImage...")
+
+            # 1. Pegue a saída do decode (tensor)
+            decoded_tensor = decoded[0]  # se 'decoded' for um dict/tuple, ajuste conforme preciso
+
+            # 2. Instancia o SaveImage
+            saveimage_node = NODE_CLASS_MAPPINGS["SaveImage"]()
+
+            # 3. Usa o método save_images
+            #    'filename_prefix' é o prefixo do arquivo de saída
+            result_dict = saveimage_node.save_images(
+                filename_prefix="FluxRedux",  # ou algo dinâmico se preferir
+                images=decoded_tensor
+            )
+
+            # 4. Normalmente, o node 'save_images' retorna um dicionário contendo:
+            # {
+            #   'ui': {
+            #       'images': [
+            #           {'filename': 'FluxRedux_12345.png', 'subfolder': ''},
+            #           ...
+            #       ]
+            #   },
+            #   ...
+            # }
+            # Assim, para pegar o nome do arquivo salvo:
+            saved_path = os.path.join(output_dir, result_dict["ui"]["images"][0]["filename"])
+
+            logger.info(f"Imagem salva em: {saved_path}")
+            return saved_path
 
     except Exception as e:
         logger.error(f"Erro ao gerar imagem: {str(e)}")
@@ -361,5 +379,4 @@ with gr.Blocks() as app:
     )
 
 if __name__ == "__main__":
-    # Ajuste caso queira compartilhar publicamente, exemplo: app.launch(server_name="0.0.0.0", share=True)
     app.launch()