Update app.py

app.py

@@ -1,352 +1,267 @@
 import os
 import random
-import sys
 import torch
-import gradio as gr
 from pathlib import Path
+from PIL import Image
+import gradio as gr
 from huggingface_hub import hf_hub_download
-import spaces
-from typing import Union, Sequence, Mapping, Any
-from comfy import model_management
 from nodes import NODE_CLASS_MAPPINGS
-
-# 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 init_extra_nodes
 
-# 3. Configuração de Diretórios
+# Configure base and output directories
 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
-def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
-    try:
-        return obj[index]
-    except KeyError:
-        return obj["result"][index]
-
-def find_path(name: str, path: str = None) -> str:
-    if path is None:
-        path = os.getcwd()
-    if name in os.listdir(path):
-        path_name = os.path.join(path, name)
-        print(f"{name} found: {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:
-    comfyui_path = find_path("ComfyUI")
-    if comfyui_path is not None and os.path.isdir(comfyui_path):
-        sys.path.append(comfyui_path)
-        print(f"'{comfyui_path}' added to sys.path")
-
-def add_extra_model_paths() -> None:
-    try:
-        from main import load_extra_path_config
-    except ImportError:
-        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("Could not find the extra_model_paths config file.")
-
-# 7. Inicialização de caminhos
-add_comfyui_directory_toSyspath()
-add_extra_model_paths()
+# Download models
+def download_models():
+    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"),
+        ("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")
+    ]
+    
+    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"Downloading {filename} from {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)
 
-def import_custom_nodes() -> None:
+# Load custom nodes
+def import_custom_nodes():
     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()
 
-# 8. 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")
-    ]
+# 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()
     
-    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
+    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"
+            )
 
-# 9. Download e Inicialização dos Modelos
-print("Baixando modelos...")
-download_models()
+            # Text Encoding
+            cliptextencode = NODE_CLASS_MAPPINGS["CLIPTextEncode"]()
+            encoded_text = cliptextencode.encode(
+                text=prompt,
+                clip=dualcliploader_loaded[0]
+            )
 
-print("Inicializando modelos...")
-import_custom_nodes()
+            # Load Style Model
+            stylemodelloader = NODE_CLASS_MAPPINGS["StyleModelLoader"]()
+            style_model = stylemodelloader.load_style_model(
+                style_model_name="flux1-redux-dev.safetensors"
+            )
 
-# Global variables for preloaded models and constants
-intconstant = NODE_CLASS_MAPPINGS["INTConstant"]()
-CONST_1024 = intconstant.get_value(value=1024)
+            # Load CLIP Vision
+            clipvisionloader = NODE_CLASS_MAPPINGS["CLIPVisionLoader"]()
+            clip_vision = clipvisionloader.load_clip(
+                clip_name="model.safetensors"
+            )
 
-# Load CLIP
-dualcliploader = NODE_CLASS_MAPPINGS["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"
-)
+            # Load Input Image
+            loadimage = NODE_CLASS_MAPPINGS["LoadImage"]()
+            loaded_image = loadimage.load_image(image=input_image)
 
-# Load VAE
-vaeloader = NODE_CLASS_MAPPINGS["VAELoader"]()
-VAE_MODEL = vaeloader.load_vae(
-    vae_name="ae.safetensors"
-)
+            # Load VAE
+            vaeloader = NODE_CLASS_MAPPINGS["VAELoader"]()
+            vae = vaeloader.load_vae(vae_name="ae.safetensors")
 
-# Load CLIP Vision
-clipvisionloader = NODE_CLASS_MAPPINGS["CLIPVisionLoader"]()
-CLIP_VISION_MODEL = clipvisionloader.load_clip(
-    clip_name="model.safetensors"
-)
+            # Load UNET
+            unetloader = NODE_CLASS_MAPPINGS["UNETLoader"]()
+            unet = unetloader.load_unet(
+                unet_name="flux1-dev.sft",
+                weight_dtype="fp8_e4m3fn"
+            )
 
-# Load Style Model
-stylemodelloader = NODE_CLASS_MAPPINGS["StyleModelLoader"]()
-STYLE_MODEL = stylemodelloader.load_style_model(
-    style_model_name="flux1-redux-dev.safetensors"
-)
+            # 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]
+            )
 
-# Initialize samplers
-ksamplerselect = NODE_CLASS_MAPPINGS["KSamplerSelect"]()
-SAMPLER = ksamplerselect.get_sampler(sampler_name="euler")
+            # Flux Guidance
+            fluxguidance = NODE_CLASS_MAPPINGS["FluxGuidance"]()
+            flux_guidance = fluxguidance.append(
+                guidance=guidance,
+                conditioning=encoded_text[0]
+            )
 
-# Initialize other nodes
-cliptextencode = NODE_CLASS_MAPPINGS["CLIPTextEncode"]()
-loadimage = NODE_CLASS_MAPPINGS["LoadImage"]()
-vaeencode = NODE_CLASS_MAPPINGS["VAEEncode"]()
-fluxguidance = NODE_CLASS_MAPPINGS["FluxGuidance"]()
-instructpixtopixconditioning = NODE_CLASS_MAPPINGS["InstructPixToPixConditioning"]()
-clipvisionencode = NODE_CLASS_MAPPINGS["CLIPVisionEncode"]()
-stylemodelapplyadvanced = NODE_CLASS_MAPPINGS["StyleModelApplyAdvanced"]()
-emptylatentimage = NODE_CLASS_MAPPINGS["EmptyLatentImage"]()
-basicguider = NODE_CLASS_MAPPINGS["BasicGuider"]()
-basicscheduler = NODE_CLASS_MAPPINGS["BasicScheduler"]()
-randomnoise = NODE_CLASS_MAPPINGS["RandomNoise"]()
-samplerCustomAdvanced = NODE_CLASS_MAPPINGS["SamplerCustomAdvanced"]()
-vaedecode = NODE_CLASS_MAPPINGS["VAEDecode"]()
-saveimage = NODE_CLASS_MAPPINGS["SaveImage"]()
-getimagesizeandcount = NODE_CLASS_MAPPINGS["GetImageSizeAnd Count"]()
-depthanything_v2 = NODE_CLASS MAPPINGS["DepthAnything_V2"]()
-cr_text = NODE_CLASS_MAPPINGS["CR Text"]()
+            # 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]
+            )
 
-model_loaders = [CLIP_MODEL, VAE_MODEL, CLIP_VISION_MODEL, STYLE_MODEL]
+            # Empty Latent Image
+            emptylatentimage = NODE_CLASS_MAPPINGS["EmptyLatentImage"]()
+            empty_latent = emptylatentimage.generate(
+                width=width,
+                height=height,
+                batch_size=batch_size
+            )
 
-model_management.load_models_gpu([
-    loader[0].patcher if hasattr(loader[0], 'patcher') else loader[0] for loader in model_loaders
-])
+            # 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]
+            )
 
-@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)) -> str:
-    with torch.inference_mode():
-        # Set up CLIP
-        clip_switch = cr_text.text_multiline(text="Flux_BFL_Depth_Redux")
-        
-        # Encode text
-        text_encoded = cliptextencode.encode(
-            text=prompt,
-            clip=get_value_at_index(CLIP_MODEL, 0),
-        )
-        
-        # Process input image
-        loaded_image = loadimage.load_image(image=image=input_image)
-        
-        # Get image size
-        size_info = getimagesizeandcount.getsize(
-            image=get_value_at_index(loaded_image, 0)
-        )
-        
-        # Encode VAE
-        vae_encoded = vaeencode.encode(
-            pixels=get_value_at_index(size_info, 0),
-            vae=get_value_at_index(Vae_model, 0),
-        )
-        
-        # Apply Flux guidance
-        flux guided = flux Guidance.append(
-            guidance=guidance,
-            conditioning=get_valueAtIndex(text_encoded, 0),
-        )
-        
-        # Set up empty latent
-        empty_latent = empty_latentimage.generate(
-            width=width,
-            height=height,
-            batch_size=batch_size
-        )
-        
-        # Set up guidance
-        guided = basicguider.get_guider(
-            model=get_value_at_index(unet_model, 0),
-            conditioning=get_value_at_index(loaded_image, 0)
-        )
-        
-        # Set up scheduler
-        schedule = basicscheduler.get_sigmas(
-            scheduler="simple",
-            steps=steps,
-            denoise=1,
-            model=get_value_atIndex(Unet Model, 0),
-        )
-        
-        # Generate random noise
-        noise = randomnoise.get_noise(noise_seed=seed)
-        
-        # Sample
-        sampled = samplerCustom advanced.sample(
-            noise=get_value_at_index(noise, 0),
-            guider=get_value at Index(guided, 0),
-            sampler=get_value at index(sampler, 0),
-            sigmas=get_value at Index(schedule, 0),
-            latent_image=get_value_atindex(empty_latent, 0)
-        )
-        
-        # Decode VAE
-        decoded = va edecode.decode(
-            samples=get_value_atindex(sampled, 0),
-            vae=get_value_at Index(VAE Model, 0),
-        )
-        
-        # Save image
-        saved = saveimage.save_images(
-            filename_prefix=get_value at index(clip switch, 0),
-            images=getValueAtIndex(decoded, 0),
-        )
-        
-        saved_path = f"output/{saved['ui']['images'][0]['filename']}"
-        
-        return saved_path
+            # 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)
 
-# Create Gradio interface
-examples = [
-    ["", "mona.png", 0.5, 3.5, 3, 1.0, random.randint(1, 2**64), 1024, 1024, 1, 20],
-    ["a woman looking at a house catching fire on the background", "disaster Girl.png", 0.6, 3.5, 3, 1.0, random.randint(1, 2**64), 1024, 1024, 1, 20],
-    ["Istanbul aerial, dramatic photography", "Natasha.png", 0.5, 3.5, 3, 1.0, random.randint(1, 2**64), 1024, 1024, 1, 20],
-]
+            # 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}")
 
-output_image = gr.Image(label="Generated image")
+            # 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")
+    gr.Markdown("# FLUX Redux Image Generator")
     
     with gr.Row():
-        with gr.column():
-            prompt_input = gr.Text box(
+        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(
+            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(
+                    guidance = gr.Slider(
                         minimum=0,
                         maximum=20,
                         step=0.1,
                         value=3.5,
                         label="Guidance"
                     )
-                    downsampling_factor = gr.slider(
-                        minimum=0,
+                    downsampling_factor = gr.Slider(
+                        minimum=1,
                         maximum=8,
                         step=1,
                         value=3,
-                        label="Downsampling factor"
+                        label="Downsampling Factor"
                     )
-                    weight = gr.slider(
+                    weight = gr.Slider(
                         minimum=0,
                         maximum=2,
                         step=0.1,
                         value=1.0,
-                        label="Model weight"
+                        label="Model Weight"
                     )
-                    seed = gr.number(
+                with gr.Column():
+                    seed = gr.Number(
                         value=random.randint(1, 2**64),
-                        label="seed",
+                        label="Seed",
                         precision=0
                     )
-                    width = gr.number(
+                    width = gr.Number(
                         value=1024,
-                        label="width",
+                        label="Width",
                         precision=0
                     )
-                    height = gr.number(
+                    height = gr.Number(
                         value=1024,
-                        label="height",
+                        label="Height",
                         precision=0
                     )
-                    batch_size = gr.number(
+                    batch_size = gr.Number(
                         value=1,
-                        label="batch size",
+                        label="Batch Size",
                         precision=0
                     )
-                    steps = gr.number(
+                    steps = gr.Number(
                         value=20,
-                        label="steps",
+                        label="Steps",
                         precision=0
                     )
-                
-                with gr.column():
-                    input_image = gr Image(
-                        label="Input Image",
-                        type="filepath"
-                    )
             
-            generate_btn = gr.button("Generate image")
+            generate_btn = gr.Button("Generate Image")
         
-        with gr.column():
-            output_image.render()
-
+        with gr.Column():
+            output_image = gr.Image(label="Generated Image", type="pil")
+    
     generate_btn.click(
         fn=generate_image,
         inputs=[
@@ -366,4 +281,6 @@ with gr.Blocks() as app:
     )
 
 if __name__ == "__main__":
+    # Download models if they don't exist
+    download_models()
     app.launch(share=True)