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()