Update app.py

app.py

@@ -1,269 +1,159 @@
-import os
-import random
-import torch
-from pathlib import Path
-from PIL import Image
-import gradio as gr
-from nodes import NODE_CLASS_MAPPINGS
-import folder_paths
-
-# Configure base and output directories
-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)
-
-def import_custom_nodes():
-    """Loads custom nodes required for the workflow."""
-    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()
-
-def generate_image(prompt, input_image, lora_weight, guidance, downsampling_factor, weight, seed, width, height, batch_size, steps):
-    """
-    Main function to execute the workflow and generate an image.
-    """
-    import_custom_nodes()
-    
-    try:
-        with torch.inference_mode():
-            # 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="default"
-            )
-
-            # Text Encoding
-            cliptextencode = NODE_CLASS_MAPPINGS["CLIPTextEncode"]()
-            encoded_text = cliptextencode.encode(
-                text=prompt,
-                clip=dualcliploader_loaded[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="sigclip_vision_patch14_384.safetensors"
-            )
-
-            # Load Input Image
-            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(
-                guidance=guidance,
-                conditioning=encoded_text[0]
-            )
-
-            # 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,
-                autocrop_margin=0.1,
-                conditioning=flux_guidance[0],
-                style_model=style_model[0],
-                clip_vision=clip_vision[0],
-                image=loaded_image[0]
-            )
-
-            # Empty Latent Image
-            emptylatentimage = NODE_CLASS_MAPPINGS["EmptyLatentImage"]()
-            empty_latent = emptylatentimage.generate(
-                width=width,
-                height=height,
-                batch_size=batch_size
-            )
-
-            # KSampler
-            ksampler = NODE_CLASS_MAPPINGS["KSampler"]()
-            sampled = ksampler.sample(
-                seed=seed,
-                steps=steps,
-                cfg=1,
-                sampler_name="euler",
-                scheduler="simple",
-                denoise=1,
-                model=lora_model[0],
-                positive=redux_result[0],
-                negative=flux_guidance[0],
-                latent_image=empty_latent[0]
-            )
-
-            # VAE Decode
-            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]
-            )
-
-            # 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
-
-    except Exception as e:
-        print(f"Error during generation: {str(e)}")
-        return None
-
-# Gradio Interface
-with gr.Blocks() as app:
-    gr.Markdown("# FLUX Redux Image Generator")
-    
-    with gr.Row():
-        with gr.Column():
-            prompt_input = gr.Textbox(
-                label="Prompt",
-                placeholder="Enter your prompt here...",
-                lines=5
-            )
-            input_image = gr.Image(
-                label="Input Image",
-                type="filepath"
-            )
-            
-            with gr.Row():
-                with gr.Column():
-                    lora_weight = gr.Slider(
-                        minimum=0,
-                        maximum=2,
-                        step=0.1,
-                        value=0.6,
-                        label="LoRA Weight"
-                    )
-                    guidance = gr.Slider(
-                        minimum=0,
-                        maximum=20,
-                        step=0.1,
-                        value=3.5,
-                        label="Guidance"
-                    )
-                    downsampling_factor = gr.Slider(
-                        minimum=1,
-                        maximum=8,
-                        step=1,
-                        value=3,
-                        label="Downsampling Factor"
-                    )
-                    weight = gr.Slider(
-                        minimum=0,
-                        maximum=2,
-                        step=0.1,
-                        value=1.0,
-                        label="Model Weight"
-                    )
-                with gr.Column():
-                    seed = gr.Number(
-                        value=random.randint(1, 2**64),
-                        label="Seed",
-                        precision=0
-                    )
-                    width = gr.Number(
-                        value=1024,
-                        label="Width",
-                        precision=0
-                    )
-                    height = gr.Number(
-                        value=1024,
-                        label="Height",
-                        precision=0
-                    )
-                    batch_size = gr.Number(
-                        value=1,
-                        label="Batch Size",
-                        precision=0
-                    )
-                    steps = gr.Number(
-                        value=20,
-                        label="Steps",
-                        precision=0
-                    )
-            
-            generate_btn = gr.Button("Generate Image")
-        
-        with gr.Column():
-            output_image = gr.Image(label="Generated Image", type="pil")
-    
-    generate_btn.click(
-        fn=generate_image,
-        inputs=[
-            prompt_input,
-            input_image,
-            lora_weight,
-            guidance,
-            downsampling_factor,
-            weight,
-            seed,
-            width,
-            height,
-            batch_size,
-            steps
-        ],
-        outputs=[output_image]
-    )
-
-if __name__ == "__main__":
-    app.launch()
-
-
-#python app.py
+import os
+import random
+import torch
+from pathlib import Path
+from PIL import Image
+import gradio as gr
+from huggingface_hub import hf_hub_download
+from nodes import NODE_CLASS_MAPPINGS
+import folder_paths
+
+# Diretório base e de saída
+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)
+
+# Baixar os modelos necessários
+
+hf_hub_download(repo_id="black-forest-labs/FLUX.1-Redux-dev",
+filename="flux1-redux-dev.safetensors",
+local_dir="models/style_models")
+
+hf_hub_download(repo_id="comfyanonymous/flux_text_encoders",
+filename="t5xxl_fp16.safetensors",
+local_dir="models/text_encoders")
+
+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")
+
+hf_hub_download(repo_id="black-forest-labs/FLUX.1-dev",
+filename="ae.safetensors",
+local_dir="models/vae")
+
+hf_hub_download(repo_id="black-forest-labs/FLUX.1-dev",
+filename="flux1-dev.safetensors.safetensors",
+local_dir="models/diffusion_models")
+
+hf_hub_download(repo_id="google/siglip-so400m-patch14-384",
+filename="model.safetensors",
+local_dir="models/clip_vision")
+
+hf_hub_download(repo_id="nftnik/NFTNIK-FLUX.1-dev-LoRA",
+filename="NFTNIK_FLUX.1[dev]_LoRA.safetensors",
+local_dir="models/lora")
+
+# Função para importar nodes personalizados
+def import_custom_nodes():
+    """Carregar nodes customizados."""
+    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()
+
+# Função principal de geração
+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():
+            # Carregar CLIP
+            dualcliploader = NODE_CLASS_MAPPINGS["DualCLIPLoader"]()
+            dualcliploader_loaded = dualcliploader.load_clip(
+                clip_name1="models/text_encoders/t5xxl_fp16.safetensors",
+                clip_name2="models/clip_vision/ViT-L-14-TEXT-detail-improved-hiT-GmP-HF.safetensors",
+                type="flux"
+            )
+
+            # Codificar texto
+            cliptextencode = NODE_CLASS_MAPPINGS["CLIPTextEncode"]()
+            encoded_text = cliptextencode.encode(
+                text=prompt,
+                clip=dualcliploader_loaded[0]
+            )
+
+            # Carregar modelos de estilo e LoRA
+            stylemodelloader = NODE_CLASS_MAPPINGS["StyleModelLoader"]()
+            style_model = stylemodelloader.load_style_model(
+                style_model_name="models/style_models/flux1-redux-dev.safetensors"
+            )
+            loraloadermodelonly = NODE_CLASS_MAPPINGS["LoraLoaderModelOnly"]()
+            lora_model = loraloadermodelonly.load_lora_model_only(
+                lora_name="models/lora/NFTNIK_FLUX.1[dev]_LoRA.safetensors",
+                strength_model=lora_weight,
+                model=style_model[0]
+            )
+
+            # Processar imagem de entrada
+            loadimage = NODE_CLASS_MAPPINGS["LoadImage"]()
+            loaded_image = loadimage.load_image(image=input_image)
+
+            # Configurações adicionais e saída
+            vaeloader = NODE_CLASS_MAPPINGS["VAELoader"]()
+            vae = vaeloader.load_vae(vae_name="models/vae/ae.safetensors")
+
+            # Decodificar e salvar
+            vaedecode = NODE_CLASS_MAPPINGS["VAEDecode"]()
+            decoded = vaedecode.decode(
+                samples=lora_model[0],
+                vae=vae[0]
+            )
+
+            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"Erro ao gerar imagem: {str(e)}")
+        return None
+
+# Interface Gradio
+with gr.Blocks() as app:
+    gr.Markdown("# Gerador de Imagens FLUX Redux")
+    with gr.Row():
+        with gr.Column():
+            prompt_input = gr.Textbox(label="Prompt", placeholder="Digite seu prompt aqui...", lines=5)
+            input_image = gr.Image(label="Imagem de Entrada", type="filepath")
+            lora_weight = gr.Slider(minimum=0, maximum=2, step=0.1, value=0.6, label="Peso LoRA")
+            guidance = gr.Slider(minimum=0, maximum=20, step=0.1, value=3.5, label="Orientação")
+            downsampling_factor = gr.Slider(minimum=1, maximum=8, step=1, value=3, label="Fator de Redução")
+            weight = gr.Slider(minimum=0, maximum=2, step=0.1, value=1.0, label="Peso do Modelo")
+            seed = gr.Number(value=random.randint(1, 2**64), label="Seed", precision=0)
+            width = gr.Number(value=1024, label="Largura", precision=0)
+            height = gr.Number(value=1024, label="Altura", precision=0)
+            batch_size = gr.Number(value=1, label="Tamanho do Lote", precision=0)
+            steps = gr.Number(value=20, label="Etapas", precision=0)
+            generate_btn = gr.Button("Gerar Imagem")
+
+        with gr.Column():
+            output_image = gr.Image(label="Imagem Gerada", type="filepath")
+
+    generate_btn.click(
+        fn=generate_image,
+        inputs=[
+            prompt_input,
+            input_image,
+            lora_weight,
+            guidance,
+            downsampling_factor,
+            weight,
+            seed,
+            width,
+            height,
+            batch_size,
+            steps
+        ],
+        outputs=[output_image]
+    )
+
+if __name__ == "__main__":
+    app.launch()