import gradio as gr
from diffusers import StableDiffusionPipeline, DDIMScheduler
import torch
from PIL import ImageEnhance, Image
import numpy as np
import spaces  # For using ZeroGPU decorator

# Load Stable Diffusion pipeline with efficient defaults
model_id = "CompVis/stable-diffusion-v1-4"
default_scheduler = DDIMScheduler.from_pretrained(model_id, subfolder="scheduler")

device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if device == "cuda" else torch.float32
pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=default_scheduler, torch_dtype=torch_dtype).to(device)

# Scheduler options (setting DDIM as the default for speed)
schedulers = {
    "High-Definition & Fast (DDIM) - Good quality with fastest speed": DDIMScheduler,
    "Photo-Realistic (PNDM) - Best for realistic details, moderate speed": PNDMScheduler,
    "Artistic & Imaginative (Euler Ancestral) - Good for creative scenes, moderate speed": EulerAncestralDiscreteScheduler,
}

# Main image generation function with 60-second GPU duration limit
@spaces.GPU(duration=60)
def generate_image(prompt, use_categories, genre, style, theme, lighting, scheduler_choice, quality, size):
    # Check if additional categories should be added to the prompt
    if use_categories:
        prompt_text = (
            f"{prompt.strip()} in a {genre.lower()} wallpaper style, "
            f"with {style.lower()} visuals, focusing on a {theme.lower()} theme "
            f"and {lighting.lower()} lighting."
        )
    else:
        prompt_text = prompt.strip()

    # Set the scheduler based on user choice
    scheduler = schedulers[scheduler_choice].from_pretrained(model_id, subfolder="scheduler")
    pipe.scheduler = scheduler

    # Set output size based on selection
    image_size = (512, 512) if size == "Profile Picture" else (512, 512)  # Using smaller size for faster generation

    # Generate image with specified quality and size
    with torch.no_grad():
        image = pipe(prompt_text, num_inference_steps=quality, guidance_scale=7.5).images[0]
        image = image.resize(image_size)  # Resize image to fit selected dimensions

    return np.array(image)  # Return as NumPy array for Gradio

# Post-processing function for brightness and contrast
def adjust_brightness_contrast(image, brightness, contrast):
    image = Image.fromarray(image.astype('uint8'), 'RGB')
    image = ImageEnhance.Brightness(image).enhance(brightness)
    image = ImageEnhance.Contrast(image).enhance(contrast)
    return np.array(image)

# Warning function for high-quality settings
def show_warning(quality):
    if quality > 40:  # Lower threshold for warning
        return "⚠️ High Quality: This may slow down generation. Consider using values below 40 for faster results."
    return ""

# Build Gradio Interface with adjusted defaults
with gr.Blocks() as demo:
    gr.Markdown("# ✨ AI-Powered Wallpaper/Profile Picture Generator\n🖼️ A tool to generate and fine-tune AI-created wallpapers and profile pictures with adjustable styles and effects.")
    
    # Image Generation Section
    with gr.Tab("Image Generator"):
        gr.Markdown("## Generate an Image")
        
        with gr.Row():
            with gr.Column():
                custom_prompt = gr.Textbox(label="Custom Prompt", placeholder="Describe your image (e.g., 'A forest at sunset')")
                
                # Toggle for using additional categories
                use_categories = gr.Checkbox(label="Enable Advanced Settings (Genre, Style, Theme, Lighting)", value=False)
                
                # Additional categories, hidden by default and shown only if use_categories is checked
                with gr.Accordion("Advanced Settings", open=False):
                    genre = gr.Dropdown(["Futuristic", "Nature", "Abstract", "Fantasy", "Sci-Fi", "Cyberpunk"], label="Genre")
                    style = gr.Dropdown(["Realistic", "Surreal", "Digital Art", "Cartoon", "Photorealistic"], label="Style")
                    theme = gr.Dropdown(["Landscape", "Portrait", "Abstract Patterns", "Architecture"], label="Theme")
                    lighting = gr.Dropdown(["Warm", "Cool", "Cinematic", "Soft", "Neon"], label="Lighting")

                # Reduced quality for faster generation
                quality = gr.Slider(20, 80, value=30, step=10, label="Image Quality", info="Lower values for faster generation.")
                warning_message = gr.Markdown("")
                
                # Scheduler selection with default option as DDIM for speed
                scheduler_choice = gr.Dropdown(
                    [
                    "High-Definition & Fast (DDIM) - Good quality with fastest speed",
                    "Photo-Realistic (PNDM) - Best for realistic details, moderate speed",
                    "Artistic & Imaginative (Euler Ancestral) - Good for creative scenes, moderate speed"
                    ],
                    label="Artistic Style & Speed",
                    value="High-Definition & Fast (DDIM) - Good quality with fastest speed"
                )
                
                size = gr.Dropdown(["Profile Picture", "Wallpaper"], label="Image Size", value="Profile Picture")
                generate_button = gr.Button("Generate Image")

            with gr.Column():
                generated_image = gr.Image(label="Generated Image", interactive=False)

        # Display warning message for high-quality settings
        quality.change(show_warning, inputs=[quality], outputs=warning_message)

        # Bind the generate function to the generate button
        generate_button.click(
            fn=generate_image,
            inputs=[custom_prompt, use_categories, genre, style, theme, lighting, scheduler_choice, quality, size],
            outputs=generated_image
        )

    # Post-Processing Section
    with gr.Tab("Edit Generated Image"):
        gr.Markdown("## Adjust Brightness & Contrast")
        
        with gr.Row():
            with gr.Column():
                brightness_slider = gr.Slider(0.5, 2.0, value=1.0, label="Brightness")
                contrast_slider = gr.Slider(0.5, 2.0, value=1.0, label="Contrast")
                apply_adjustments = gr.Button("Apply Adjustments")

            with gr.Column():
                output_image = gr.Image(label="Adjusted Image", interactive=False)

        # Bind the brightness and contrast adjustment function to the Apply Adjustments button
        apply_adjustments.click(
            fn=adjust_brightness_contrast,
            inputs=[generated_image, brightness_slider, contrast_slider],
            outputs=output_image
        )

# Launch with a public shareable link
demo.launch(share=True)