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