Update app.py

app.py

@@ -1,198 +1,321 @@
 import os
 import sys
 import random
+from typing import Sequence, Mapping, Any, Union
+
 import torch
-from pathlib import Path
-from PIL import Image
 import gradio as gr
+from PIL import Image
 from huggingface_hub import hf_hub_download
-import spaces
-from typing import Union, Sequence, Mapping, Any
-
-# 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
-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
-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))
-
-# 5. Inicialização do ComfyUI
-print("Inicializando ComfyUI...")
-init_extra_nodes()
-
-# 6. Helper Functions
+
+import spaces  # Se estiver no Hugging Face Spaces. Se não, pode remover.
+
+#####################################
+# 1. Funções auxiliares de caminho e import
+#####################################
+
+def find_path(name: str, path: str = None) -> str:
+    """Busca recursivamente por uma pasta/arquivo 'name' a partir de 'path'."""
+    if path is None:
+        path = os.getcwd()
+    if name in os.listdir(path):
+        path_name = os.path.join(path, name)
+        print(f"{name} encontrado em: {path_name}")
+        return path_name
+    parent_directory = os.path.dirname(path)
+    if parent_directory == path:
+        return None
+    return find_path(name, parent_directory)
+
+def add_comfyui_directory_to_sys_path() -> None:
+    """Adiciona o diretório ComfyUI ao sys.path, caso encontrado."""
+    comfyui_path = find_path("ComfyUI")
+    if comfyui_path is not None and os.path.isdir(comfyui_path):
+        sys.path.append(comfyui_path)
+        print(f"Diretório ComfyUI adicionado ao sys.path: {comfyui_path}")
+    else:
+        print("Não foi possível encontrar o diretório ComfyUI.")
+
+def add_extra_model_paths() -> None:
+    """
+    Carrega configurações extras de caminhos de modelos, se existir 
+    um arquivo 'extra_model_paths.yaml'.
+    """
+    try:
+        from main import load_extra_path_config
+    except ImportError:
+        # Dependendo da versão do ComfyUI, pode estar em 'utils.extra_config'
+        from utils.extra_config import load_extra_path_config
+
+    extra_model_paths = find_path("extra_model_paths.yaml")
+    if extra_model_paths is not None:
+        load_extra_path_config(extra_model_paths)
+    else:
+        print("Arquivo extra_model_paths.yaml não foi encontrado.")
+
+def import_custom_nodes() -> None:
+    """
+    Executa a inicialização de nós extras e o servidor do ComfyUI (caso necessário),
+    similar ao que ocorre no segundo script.
+    """
+    import asyncio
+    import execution
+    from nodes import init_extra_nodes
+    import server
+
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    server_instance = server.PromptServer(loop)
+    execution.PromptQueue(server_instance)
+    init_extra_nodes()
+
+#####################################
+# 2. Ajustando o ambiente ComfyUI
+#####################################
+
+add_comfyui_directory_to_sys_path()
+add_extra_model_paths()
+import_custom_nodes()
+
+#####################################
+# 3. Importando nós do ComfyUI
+#####################################
+from comfy import model_management
+from nodes import (
+    NODE_CLASS_MAPPINGS,
+    DualCLIPLoader,
+    CLIPVisionLoader,
+    StyleModelLoader,
+    VAELoader,
+    CLIPTextEncode,
+    LoadImage,
+    EmptyLatentImage,
+    VAEDecode
+)
+
+#####################################
+# 4. Download de modelos (ajuste conforme sua necessidade)
+#####################################
+
+# Exemplo de downloads (ajuste conforme seus modelos):
+os.makedirs("models/text_encoders", exist_ok=True)
+os.makedirs("models/style_models", exist_ok=True)
+os.makedirs("models/diffusion_models", exist_ok=True)
+os.makedirs("models/vae", exist_ok=True)
+os.makedirs("models/clip_vision", exist_ok=True)
+
+try:
+    print("Baixando modelo Style (flux1-redux-dev.safetensors)...")
+    hf_hub_download(repo_id="black-forest-labs/FLUX.1-Redux-dev", 
+                    filename="flux1-redux-dev.safetensors",
+                    local_dir="models/style_models")
+    print("Baixando T5 (t5xxl_fp16.safetensors)...")
+    hf_hub_download(repo_id="comfyanonymous/flux_text_encoders", 
+                    filename="t5xxl_fp16.safetensors",
+                    local_dir="models/text_encoders")
+
+    print("Baixando CLIP L (ViT-L-14) ...")
+    hf_hub_download(repo_id="zer0int/CLIP-GmP-ViT-L-14",
+                    filename="ViT-L-14-TEXT-detail-improved-hiT-GmP-HF.safetensors",
+                    local_dir="models/text_encoders")
+    print("Baixando VAE (ae.safetensors)...")
+    hf_hub_download(repo_id="black-forest-labs/FLUX.1-dev",
+                    filename="ae.safetensors",
+                    local_dir="models/vae")
+    print("Baixando flux1-dev.safetensors (modelo difusão)...")
+    hf_hub_download(repo_id="black-forest-labs/FLUX.1-dev",
+                    filename="flux1-dev.safetensors",
+                    local_dir="models/diffusion_models")
+    print("Baixando CLIP Vision (model.safetensors)...")
+    hf_hub_download(repo_id="google/siglip-so400m-patch14-384",
+                    filename="model.safetensors",
+                    local_dir="models/clip_vision")
+except Exception as e:
+    print("Algum download falhou:", e)
+
+#####################################
+# 5. Carregar modelos via ComfyUI
+#####################################
+
+# Carregando CLIP (DualCLIPLoader)
+dualcliploader = DualCLIPLoader()
+clip_model = dualcliploader.load_clip(
+    clip_name1="t5xxl_fp16.safetensors",
+    clip_name2="ViT-L-14-TEXT-detail-improved-hiT-GmP-HF.safetensors",
+    type="flux"
+)
+
+# Carregando CLIP Vision
+clipvisionloader = CLIPVisionLoader()
+clip_vision_model = clipvisionloader.load_clip(
+    clip_name="model.safetensors"
+)
+
+# Carregando Style Model
+stylemodelloader = StyleModelLoader()
+style_model = stylemodelloader.load_style_model(
+    style_model_name="flux1-redux-dev.safetensors"
+)
+
+# Carregando VAE
+vaeloader = VAELoader()
+vae_model = vaeloader.load_vae(
+    vae_name="ae.safetensors"
+)
+
+# (Opcional) Se tiver um model UNet, faça UNETLoader, etc.
+
+# Opcional: Carregar para GPU
+model_management.load_models_gpu([
+    loader[0] for loader in [clip_model, clip_vision_model, style_model, vae_model]
+])
+
+#####################################
+# 6. Funções auxiliares e placeholders
+#####################################
+
 def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
+    """Retorna o 'index' de um objeto que pode ser um dict ou lista."""
     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-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)
-            # 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
-
-# 8. Download e Inicialização dos Modelos
-print("Baixando modelos...")
-download_models()
-
-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"
-    )
+#####################################
+# 7. Definir workflow simplificado
+#####################################
 
-    # CLIP Vision
-    clipvisionloader = NODE_CLASS_MAPPINGS["CLIPVisionLoader"]()
-    clip_vision = clipvisionloader.load_clip(
-        clip_name="model.safetensors"
-    )
+@spaces.GPU  # Se estiver no Hugging Face Spaces. Senão, remova.
+def generate_image(
+    prompt: str,
+    input_image_path: str,
+    lora_weight: float,
+    guidance: float,
+    downsampling_factor: float,
+    weight: float,
+    seed: int,
+    width: int,
+    height: int,
+    batch_size: int,
+    steps: int,
+    progress=gr.Progress(track_tqdm=True)
+):
+    """
+    Gera imagem usando um fluxo simplificado, similar ao primeiro script.
+    """
+    try:
+        # Garantindo repetibilidade do seed
+        torch.manual_seed(seed)
+        random.seed(seed)
 
-    # Style Model
-    stylemodelloader = NODE_CLASS_MAPPINGS["StyleModelLoader"]()
-    stylemodelloader_441 = stylemodelloader.load_style_model(
-        style_model_name="flux1-redux-dev.safetensors"
-    )
+        # 1) Encode Texto
+        cliptextencode = CLIPTextEncode()
+        encoded_text = cliptextencode.encode(
+            text=prompt,
+            clip=get_value_at_index(clip_model, 0)
+        )
 
-    # VAE
-    vaeloader = NODE_CLASS_MAPPINGS["VAELoader"]()
-    vaeloader_359 = vaeloader.load_vae(
-        vae_name="ae.safetensors"
-    )
+        # 2) Carregar imagem de entrada
+        loadimage = LoadImage()
+        loaded_image = loadimage.load_image(image=input_image_path)
 
-# 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_357[0]
-            )
+        # 3) Flux Guidance (se existir)
+        fluxguidance = NODE_CLASS_MAPPINGS["FluxGuidance"]()
+        flux_guided = fluxguidance.append(
+            guidance=guidance,
+            conditioning=get_value_at_index(encoded_text, 0)
+        )
 
-            # Carregar e processar imagem
-            loadimage = NODE_CLASS_MAPPINGS["LoadImage"]()
-            loaded_image = loadimage.load_image(image=input_image)
+        # 4) Redux Advanced (aplicar style model)
+        reduxadvanced = NODE_CLASS_MAPPINGS["ReduxAdvanced"]()
+        redux_result = reduxadvanced.apply_stylemodel(
+            downsampling_factor=downsampling_factor,
+            downsampling_function="area",
+            mode="keep aspect ratio",
+            weight=weight,
+            conditioning=get_value_at_index(flux_guided, 0),
+            style_model=get_value_at_index(style_model, 0),
+            clip_vision=get_value_at_index(clip_vision_model, 0),
+            image=get_value_at_index(loaded_image, 0)
+        )
 
-            # Flux Guidance
-            fluxguidance = NODE_CLASS_MAPPINGS["FluxGuidance"]()
-            flux_guidance = fluxguidance.append(
-                guidance=guidance,
-                conditioning=encoded_text[0]
-            )
+        # 5) Empty Latent
+        emptylatent = EmptyLatentImage()
+        empty_latent = emptylatent.generate(
+            width=width,
+            height=height,
+            batch_size=batch_size
+        )
 
-            # Redux Advanced
-            reduxadvanced = NODE_CLASS_MAPPINGS["ReduxAdvanced"]()
-            redux_result = reduxadvanced.apply_stylemodel(
-                downsampling_factor=downsampling_factor,
-                downsampling_function="area",
-                mode="keep aspect ratio",
-                weight=weight,
-                conditioning=flux_guidance[0],
-                style_model=stylemodelloader_441[0],
-                clip_vision=clip_vision[0],
-                image=loaded_image[0]
-            )
+        # 6) KSampler (no ComfyUI atual, há "KSamplerSelect" ou "KSampler")
+        ksampler = NODE_CLASS_MAPPINGS["KSampler"]()
+        sampled = ksampler.sample(
+            seed=seed,
+            steps=steps,
+            cfg=1,  # Exemplo de CFG = 1
+            sampler_name="euler",
+            scheduler="simple",
+            denoise=1,
+            model=get_value_at_index(style_model, 0),     # Usa o style model como UNet? (depende da config)
+            positive=get_value_at_index(redux_result, 0),
+            negative=get_value_at_index(flux_guided, 0),
+            latent_image=get_value_at_index(empty_latent, 0)
+        )
 
-            # Empty Latent
-            emptylatentimage = NODE_CLASS_MAPPINGS["EmptyLatentImage"]()
-            empty_latent = emptylatentimage.generate(
-                width=width,
-                height=height,
-                batch_size=batch_size
-            )
+        # 7) Decodificar VAE
+        vaedecode = VAEDecode()
+        decoded = vaedecode.decode(
+            samples=get_value_at_index(sampled, 0),
+            vae=get_value_at_index(vae_model, 0)
+        )
 
-            # KSampler
-            ksampler = NODE_CLASS_MAPPINGS["KSampler"]()
-            sampled = ksampler.sample(
-                seed=seed,
-                steps=steps,
-                cfg=1,
-                sampler_name="euler",
-                scheduler="simple",
-                denoise=1,
-                model=stylemodelloader_441[0],
-                positive=redux_result[0],
-                negative=flux_guidance[0],
-                latent_image=empty_latent[0]
-            )
+        # 8) Salvar imagem
+        output_dir = "output"
+        os.makedirs(output_dir, exist_ok=True)
+        temp_filename = f"Flux_{random.randint(0, 99999)}.png"
+        temp_path = os.path.join(output_dir, temp_filename)
 
-            # Decodificar VAE
-            vaedecode = NODE_CLASS_MAPPINGS["VAEDecode"]()
-            decoded = vaedecode.decode(
-                samples=sampled[0],
-                vae=vaeloader_359[0]
-            )
+        # No ComfyUI, 'decoded[0]' pode ser um tensor [C,H,W] normalizado
+        # ou algo no formato [N,C,H,W]. Precisamos converter para PIL:
+        # Se for um batch, pegue o primeiro item. Ajuste se quiser batch maior.
+        image_data = get_value_at_index(decoded, 0)
+        # Normalmente, se for "float [0,1]" em C,H,W:
+        # Precisamos mover pro CPU e converter em numpy
+        if isinstance(image_data, torch.Tensor):
+            image_data = image_data.cpu().numpy()
+
+        # Se a imagem estiver em [C,H,W], transpor para [H,W,C] e escalar 0..255
+        if len(image_data.shape) == 3:
+            image_data = image_data.transpose(1, 2, 0)
+        image_data = (image_data * 255).clip(0, 255).astype("uint8")
 
-            # 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)
+        pil_image = Image.fromarray(image_data)
+        pil_image.save(temp_path)
 
-            return temp_path
+        return temp_path
     except Exception as e:
         print(f"Erro ao gerar imagem: {str(e)}")
         return None
 
-# 10. Interface Gradio
+#####################################
+# 8. Interface Gradio (similar ao primeiro snippet)
+#####################################
+
 with gr.Blocks() as app:
-    gr.Markdown("# FLUX Redux Image Generator")
-    
+    gr.Markdown("# FLUX Redux Image Generator (Simplificado)")
+
     with gr.Row():
         with gr.Column():
             prompt_input = gr.Textbox(
                 label="Prompt",
-                placeholder="Enter your prompt here...",
+                placeholder="Escreva seu prompt...",
                 lines=5
             )
             input_image = gr.Image(
-                label="Input Image",
+                label="Imagem de Entrada",
                 type="filepath"
             )
-            
+
             with gr.Row():
                 with gr.Column():
                     lora_weight = gr.Slider(
@@ -200,7 +323,7 @@ with gr.Blocks() as app:
                         maximum=2,
                         step=0.1,
                         value=0.6,
-                        label="LoRA Weight"
+                        label="LoRA Weight (não usado nesse fluxo)"
                     )
                     guidance = gr.Slider(
                         minimum=0,
@@ -221,7 +344,7 @@ with gr.Blocks() as app:
                         maximum=2,
                         step=0.1,
                         value=1.0,
-                        label="Model Weight"
+                        label="Redux Model Weight"
                     )
                 with gr.Column():
                     seed = gr.Number(
@@ -230,12 +353,12 @@ with gr.Blocks() as app:
                         precision=0
                     )
                     width = gr.Number(
-                        value=1024,
+                        value=512,
                         label="Width",
                         precision=0
                     )
                     height = gr.Number(
-                        value=1024,
+                        value=512,
                         label="Height",
                         precision=0
                     )
@@ -249,12 +372,12 @@ with gr.Blocks() as app:
                         label="Steps",
                         precision=0
                     )
-            
+
             generate_btn = gr.Button("Generate Image")
-        
+
         with gr.Column():
             output_image = gr.Image(label="Generated Image", type="filepath")
-    
+
     generate_btn.click(
         fn=generate_image,
         inputs=[
@@ -274,4 +397,5 @@ with gr.Blocks() as app:
     )
 
 if __name__ == "__main__":
-    app.launch()
+    # Você pode usar app.launch(share=True) se quiser compartilhar via link.
+    app.launch()