import spaces
import argparse
import os
import time
from os import path
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download

cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
# TRANSFORMERS_CACHE is deprecated, only use HF_HOME
os.environ["HF_HUB_CACHE"] = cache_path
os.environ["HF_HOME"] = cache_path

import gradio as gr
import torch

# Try to handle version compatibility issues
try:
    from diffusers import FluxPipeline
except ImportError as e:
    print(f"Error importing FluxPipeline: {e}")
    print("Attempting to use StableDiffusionPipeline as fallback...")
    from diffusers import StableDiffusionPipeline as FluxPipeline

torch.backends.cuda.matmul.allow_tf32 = True

class timer:
    def __init__(self, method_name="timed process"):
        self.method = method_name
    def __enter__(self):
        self.start = time.time()
        print(f"{self.method} starts")
    def __exit__(self, exc_type, exc_val, exc_tb):
        end = time.time()
        print(f"{self.method} took {str(round(end - self.start, 2))}s")

if not path.exists(cache_path):
    os.makedirs(cache_path, exist_ok=True)

pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"))
pipe.fuse_lora(lora_scale=0.125)
pipe.to(device="cuda", dtype=torch.bfloat16)

# Custom CSS for gradient effects and visual enhancements
custom_css = """
.container {
    max-width: 1200px;
    margin: 0 auto;
    padding: 20px;
}

.gradio-container {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 50%, #f093fb 100%);
    min-height: 100vh;
}

.main-content {
    background: rgba(255, 255, 255, 0.95);
    border-radius: 20px;
    padding: 30px;
    box-shadow: 0 20px 40px rgba(0, 0, 0, 0.1);
    backdrop-filter: blur(10px);
}

h1 {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    text-align: center;
    font-size: 3rem !important;
    font-weight: 800 !important;
    margin-bottom: 1rem !important;
    text-shadow: 2px 2px 4px rgba(0, 0, 0, 0.1);
}

.subtitle {
    text-align: center;
    color: #666;
    font-size: 1.2rem;
    margin-bottom: 2rem;
}

.gr-button-primary {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    border: none !important;
    color: white !important;
    font-weight: bold !important;
    font-size: 1.1rem !important;
    padding: 12px 30px !important;
    border-radius: 10px !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 4px 15px rgba(102, 126, 234, 0.3) !important;
}

.gr-button-primary:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 20px rgba(102, 126, 234, 0.4) !important;
}

.gr-input, .gr-box {
    border-radius: 10px !important;
    border: 2px solid #e0e0e0 !important;
    transition: all 0.3s ease !important;
}

.gr-input:focus {
    border-color: #667eea !important;
    box-shadow: 0 0 0 3px rgba(102, 126, 234, 0.1) !important;
}

.gr-form {
    background: white !important;
    border-radius: 15px !important;
    padding: 20px !important;
    box-shadow: 0 4px 10px rgba(0, 0, 0, 0.05) !important;
}

.gr-padded {
    padding: 15px !important;
}

.badge-container {
    display: flex;
    justify-content: center;
    gap: 12px;
    margin: 20px 0;
}

.how-to-use {
    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
    border-radius: 15px;
    padding: 25px;
    margin-top: 30px;
    box-shadow: 0 4px 10px rgba(0, 0, 0, 0.05);
}

.how-to-use h2 {
    color: #667eea;
    font-size: 1.8rem;
    margin-bottom: 1rem;
}

.how-to-use ol {
    color: #555;
    line-height: 1.8;
}

.how-to-use li {
    margin-bottom: 10px;
}

.tip {
    background: rgba(102, 126, 234, 0.1);
    border-left: 4px solid #667eea;
    padding: 15px;
    margin-top: 20px;
    border-radius: 5px;
    color: #555;
    font-style: italic;
}
"""

with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
    with gr.Column(elem_classes="main-content"):
        gr.HTML(
            """
            <div style="text-align: center; max-width: 800px; margin: 0 auto;">
                <h1>FLUX Fast & Furious</h1>
                <p class="subtitle">Lightning-fast image generation powered by Hyper-FLUX LoRA</p>
            </div>
            """
        )
        
        gr.HTML(
            """
            <div class='badge-container'>
                <a href="https://huggingface.co/spaces/openfree/Best-AI" target="_blank">
                    <img src="https://img.shields.io/static/v1?label=OpenFree&message=BEST%20AI%20Services&color=%230000ff&labelColor=%23000080&logo=huggingface&logoColor=%23ffa500&style=for-the-badge" alt="OpenFree badge">
                </a>
        
                <a href="https://discord.gg/openfreeai" target="_blank">
                    <img src="https://img.shields.io/static/v1?label=Discord&message=Openfree%20AI&color=%230000ff&labelColor=%23800080&logo=discord&logoColor=white&style=for-the-badge" alt="Discord badge">
                </a>
            </div>
            """
        )

        with gr.Row():
            with gr.Column(scale=3):
                with gr.Group():
                    prompt = gr.Textbox(
                        label="✨ Your Image Description",
                        placeholder="E.g., A serene landscape with mountains and a lake at sunset",
                        lines=3
                    )
                    
                    with gr.Accordion("🎨 Advanced Settings", open=False):
                        with gr.Group():
                            with gr.Row():
                                height = gr.Slider(label="Height", minimum=256, maximum=1152, step=64, value=1024)
                                width = gr.Slider(label="Width", minimum=256, maximum=1152, step=64, value=1024)
                            
                            with gr.Row():
                                steps = gr.Slider(label="Inference Steps", minimum=6, maximum=25, step=1, value=8)
                                scales = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=5.0, step=0.1, value=3.5)
                            
                            seed = gr.Number(label="Seed (for reproducibility)", value=3413, precision=0)
                    
                    generate_btn = gr.Button("🚀 Generate Image", variant="primary", scale=1)

            with gr.Column(scale=4):
                output = gr.Image(label="🎨 Your Generated Image")
        
        gr.HTML(
            """
            <div class="how-to-use">
                <h2>📖 How to Use</h2>
                <ol>
                    <li>✍️ Enter a detailed description of the image you want to create</li>
                    <li>⚙️ Adjust advanced settings if desired (tap to expand)</li>
                    <li>🎯 Tap "Generate Image" and watch the magic happen!</li>
                </ol>
                <div class="tip">
                    💡 <strong>Pro Tip:</strong> Be specific in your description for best results! Include details about style, mood, colors, and composition.
                </div>
            </div>
            """
        )

    @spaces.GPU
    def process_image(height, width, steps, scales, prompt, seed):
        global pipe
        with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
            return pipe(
                prompt=[prompt],
                generator=torch.Generator().manual_seed(int(seed)),
                num_inference_steps=int(steps),
                guidance_scale=float(scales),
                height=int(height),
                width=int(width),
                max_sequence_length=256
            ).images[0]

    generate_btn.click(
        process_image,
        inputs=[height, width, steps, scales, prompt, seed],
        outputs=output
    )

if __name__ == "__main__":
    demo.launch()