import gradio as gr
import numpy as np
import spaces 
import torch
import random
import time
from PIL import Image
from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler, FluxTransformer2DModel
from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast, AutoProcessor, pipeline
from huggingface_hub import hf_hub_download
from gradio_client import Client, handle_file
import os
import subprocess

subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

# Use the 'waffles' environment variable as the access token
hf_token = os.getenv('waffles')

# Ensure the token is loaded correctly
if not hf_token:
    raise ValueError("Hugging Face API token not found. Please set the 'waffles' environment variable.")

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048

pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16, revision="refs/pr/1", token=hf_token).to(device)

@spaces.GPU(duration=60)
def infer(prompt, seed=0, randomize_seed=True, width=640, height=1024, guidance_scale=0.0, num_inference_steps=5, lora_model="AlekseyCalvin/RCA_Agitprop_Manufactory", progress=gr.Progress(track_tqdm=True)):
    global pipe
    
    # Load LoRA if specified
    if lora_model:
        try:
            pipe.load_lora_weights(lora_model)
        except Exception as e:
            return None, seed, f"Failed to load LoRA model: {str(e)}"

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)
    
    try:
        image = pipe(
            prompt=prompt, 
            width=width,
            height=height,
            num_inference_steps=num_inference_steps, 
            generator=generator,
            guidance_scale=guidance_scale
        ).images[0]
        
        # Unload LoRA weights after generation
        if lora_model:
            pipe.unload_lora_weights()
        
        return image, prompt, seed, "Image generated successfully."
    except Exception as e:
        return None, seed, f"Error during image generation: {str(e)}"

    
    return image, prompt, seed


examples = [
    "RCA style communist party poster with the words Ready for REVOLUTION? in large black consistent constructivist font alongside a red Soviet hammer and a red Soviet sickle over the background of planet earth, over the North American continent",
]

custom_css = """
#col-container {
    margin: 0 auto;
    max-width: 520px; 
}
.input-group, .output-group {
    border: 1px solid #eb3109;
    border-radius: 10px;
    padding: 20px;
    margin-bottom: 20px;
    background-color: #f9f9f9;
}
.submit-btn {
    background-color: #2980b9 !important;
    color: white !important;
}
.submit-btn:hover {
    background-color: #3498db !important;
}
"""

css="""
#col-container {
    margin: 0 auto;
    max-width: 520px;
}
"""

with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="red", secondary_hue="gray")) as demo:
    
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""# RCA Agitprop Manufactory: pre-phrase prompts with 'RCA style' to activate custom model """)
        
        with gr.Row():
            
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=2,
                placeholder="RCA style communist poster of ",
                container=False,
            )
            
            run_button = gr.Button("Run", scale=0)
        
        output_image = gr.Image(label="Result", elem_id="gallery", show_label=False)
        
        with gr.Accordion("Advanced Settings", open=True):
            
            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )
            
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
            
            with gr.Row():
                
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=640,
                )
                
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
            
            with gr.Row():
                
  
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=5,
                )
        
        gr.Examples(
            examples = examples,
            fn = infer,
            inputs = [prompt],
            outputs = [output_image, seed],
            cache_examples="lazy"
        )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn = infer,
        inputs = [prompt, seed, randomize_seed, width, height, num_inference_steps],
        outputs = [output_image, seed]
    )

demo.launch(debug=True)