Update app.py

app.py

@@ -9,17 +9,15 @@ 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
+# 1. Configurar logging para debug
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# 1. Configuração de Caminhos e Imports
+# 2. Configuração de Caminhos e Imports
 current_dir = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(current_dir)
 
@@ -31,7 +29,7 @@ os.makedirs(output_dir, exist_ok=True)
 os.makedirs(models_dir, exist_ok=True)
 folder_paths.set_output_directory(output_dir)
 
-# Configurar caminhos dos modelos e verificar estrutura
+# 4. Configurar caminhos dos modelos e verificar estrutura
 MODEL_FOLDERS = ["style_models", "text_encoders", "vae", "unet", "clip_vision"]
 for model_folder in MODEL_FOLDERS:
     folder_path = os.path.join(models_dir, model_folder)
@@ -39,7 +37,7 @@ for model_folder in MODEL_FOLDERS:
     folder_paths.add_model_folder_path(model_folder, folder_path)
     logger.info(f"Pasta de modelo configurada: {model_folder}")
 
-# 4. Diagnóstico CUDA
+# 5. Diagnóstico CUDA
 logger.info(f"Python version: {sys.version}")
 logger.info(f"Torch version: {torch.__version__}")
 logger.info(f"CUDA disponível: {torch.cuda.is_available()}")
@@ -47,7 +45,7 @@ logger.info(f"Quantidade de GPUs: {torch.cuda.device_count()}")
 if torch.cuda.is_available():
     logger.info(f"GPU atual: {torch.cuda.get_device_name(0)}")
 
-# 5. Inicialização do ComfyUI
+# 6. Inicialização do ComfyUI
 logger.info("Inicializando ComfyUI...")
 try:
     init_extra_nodes()
@@ -55,7 +53,7 @@ except Exception as e:
     logger.warning(f"Aviso na inicialização de nós extras: {str(e)}")
     logger.info("Continuando mesmo com avisos nos nós extras...")
 
-# 6. Helper Functions
+# 7. Helper Functions
 def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
     try:
         return obj[index]
@@ -69,7 +67,7 @@ def verify_file_exists(folder: str, filename: str) -> bool:
         logger.error(f"Arquivo não encontrado: {file_path}")
     return exists
 
-# 7. Download de Modelos
+# 8. Download de Modelos
 logger.info("Baixando modelos necessários...")
 try:
     hf_hub_download(
@@ -106,7 +104,7 @@ except Exception as e:
     logger.error(f"Erro ao baixar modelos: {str(e)}")
     raise
 
-# 8. Inicialização dos Modelos
+# 9. Inicialização dos Modelos
 logger.info("Inicializando modelos...")
 try:
     with torch.no_grad():
@@ -169,7 +167,7 @@ except Exception as e:
     logger.error(f"Erro ao inicializar modelos: {str(e)}")
     raise
 
-# 9. Função de Geração
+# 10. Função de Geração
 @spaces.GPU
 def generate_image(
     prompt, input_image, lora_weight, guidance, downsampling_factor,
@@ -246,35 +244,15 @@ def generate_image(
                 vae=VAE_MODEL[0]
             )
 
-            # ======================== SALVAR IMAGEM USANDO O NODE SaveImage ======================
+            # Salvar Imagem
             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
+            decoded_tensor = decoded[0] 
             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
+                filename_prefix="Flux_", 
                 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