import gradio as gr
import torch
import spaces
from diffusers import OvisImagePipeline

# Load the model at startup
pipe = OvisImagePipeline.from_pretrained(
    "AIDC-AI/Ovis-Image-7B", 
    torch_dtype=torch.bfloat16
)
pipe.to("cuda")

@spaces.GPU(duration=120)
def generate_image(
    prompt: str,
    progress=gr.Progress(track_tqdm=True)
):
    if not prompt.strip():
        raise gr.Error("Please enter a prompt to generate an image.")
    
    image = pipe(
        prompt=prompt,
        negative_prompt="",
        num_inference_steps=50,
        true_cfg_scale=5.0
    ).images[0]
    
    return image

# Apple-inspired CSS that works for both light and dark mode
apple_css = """
@import url('https://fonts.googleapis.com/css2?family=SF+Pro+Display:wght@300;400;500;600;700&family=SF+Pro+Text:wght@300;400;500;600&display=swap');

* {
    font-family: -apple-system, BlinkMacSystemFont, 'SF Pro Display', 'SF Pro Text', 'Helvetica Neue', sans-serif !important;
}

/* Light mode styles */
.gradio-container {
    max-width: 700px !important;
    margin: 0 auto !important;
    padding: 0 16px !important;
}

/* Main header */
.main-header {
    text-align: center;
    padding: 40px 20px 32px;
}

.main-header h1 {
    font-size: 40px;
    font-weight: 600;
    letter-spacing: -0.025em;
    margin: 0 0 8px 0;
    background: linear-gradient(135deg, #1d1d1f 0%, #424245 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
}

.main-header p {
    font-size: 17px;
    font-weight: 400;
    margin: 0;
    opacity: 0.7;
}

/* Dark mode header */
.dark .main-header h1 {
    background: linear-gradient(135deg, #f5f5f7 0%, #a1a1a6 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
}

/* Card styling */
.card-container {
    border-radius: 20px;
    padding: 24px;
    margin-bottom: 16px;
    background: rgba(255, 255, 255, 0.8);
    backdrop-filter: blur(20px);
    -webkit-backdrop-filter: blur(20px);
    border: 1px solid rgba(0, 0, 0, 0.08);
    box-shadow: 0 4px 24px rgba(0, 0, 0, 0.04);
}

.dark .card-container {
    background: rgba(30, 30, 32, 0.8);
    border: 1px solid rgba(255, 255, 255, 0.08);
    box-shadow: 0 4px 24px rgba(0, 0, 0, 0.2);
}

/* Textbox styling */
textarea {
    background: rgba(142, 142, 147, 0.12) !important;
    border: none !important;
    border-radius: 12px !important;
    padding: 16px !important;
    font-size: 16px !important;
    line-height: 1.5 !important;
    transition: all 0.2s cubic-bezier(0.4, 0, 0.2, 1) !important;
    resize: none !important;
}

textarea:focus {
    background: rgba(142, 142, 147, 0.18) !important;
    box-shadow: 0 0 0 4px rgba(0, 122, 255, 0.15) !important;
    outline: none !important;
}

textarea::placeholder {
    opacity: 0.5 !important;
}

.dark textarea {
    background: rgba(142, 142, 147, 0.16) !important;
}

.dark textarea:focus {
    background: rgba(142, 142, 147, 0.24) !important;
}

/* Label styling */
label span {
    font-size: 14px !important;
    font-weight: 500 !important;
    letter-spacing: -0.01em !important;
}

/* Button styling */
button.primary {
    background: linear-gradient(180deg, #0a84ff 0%, #0077ed 100%) !important;
    color: white !important;
    border: none !important;
    border-radius: 980px !important;
    padding: 14px 32px !important;
    font-size: 17px !important;
    font-weight: 500 !important;
    letter-spacing: -0.01em !important;
    cursor: pointer !important;
    transition: all 0.2s cubic-bezier(0.4, 0, 0.2, 1) !important;
    box-shadow: 0 2px 8px rgba(10, 132, 255, 0.3) !important;
    min-height: 50px !important;
}

button.primary:hover {
    transform: scale(1.02) !important;
    box-shadow: 0 4px 16px rgba(10, 132, 255, 0.4) !important;
}

button.primary:active {
    transform: scale(0.98) !important;
}

/* Output image styling */
.output-container {
    border-radius: 20px;
    overflow: hidden;
    background: rgba(142, 142, 147, 0.08);
    min-height: 300px;
}

.dark .output-container {
    background: rgba(142, 142, 147, 0.12);
}

/* Image container */
.output-container img {
    border-radius: 16px;
}

/* Footer styling */
.footer-link {
    text-align: center;
    padding: 24px 0;
    font-size: 13px;
}

.footer-link a {
    color: #0a84ff;
    text-decoration: none;
    font-weight: 500;
    transition: opacity 0.2s;
}

.footer-link a:hover {
    opacity: 0.7;
}

/* Remove default Gradio styling */
.block.padded:not(.gradio-group) {
    padding: 0 !important;
    background: transparent !important;
    border: none !important;
    box-shadow: none !important;
}

/* Progress bar styling */
.progress-bar {
    background: linear-gradient(90deg, #0a84ff, #5ac8fa) !important;
    border-radius: 4px !important;
}

/* Mobile responsive */
@media (max-width: 640px) {
    .gradio-container {
        padding: 0 12px !important;
    }
    
    .main-header {
        padding: 32px 16px 24px;
    }
    
    .main-header h1 {
        font-size: 32px;
    }
    
    .main-header p {
        font-size: 15px;
    }
    
    .card-container {
        padding: 20px;
        border-radius: 16px;
    }
    
    button.primary {
        width: 100% !important;
        padding: 16px 24px !important;
    }
    
    textarea {
        font-size: 16px !important;
    }
}

/* Smooth transitions for dark mode */
.gradio-container,
.card-container,
textarea,
button {
    transition: background-color 0.3s ease, border-color 0.3s ease, box-shadow 0.3s ease !important;
}
"""

with gr.Blocks(title="OVIS Image Generator") as demo:
    
    # Header
    gr.HTML("""
        <div class="main-header">
            <h1>OVIS Image</h1>
            <p>Transform your ideas into images</p>
        </div>
    """)
    
    # Input card
    with gr.Column(elem_classes="card-container"):
        prompt = gr.Textbox(
            label="Describe your image",
            placeholder="A serene mountain landscape at golden hour with snow-capped peaks...",
            lines=3,
            show_label=True,
            max_lines=5
        )
        
        generate_btn = gr.Button(
            "Generate",
            variant="primary",
            size="lg"
        )
    
    # Output card
    with gr.Column(elem_classes="card-container output-container"):
        output_image = gr.Image(
            label="Generated Image",
            type="pil",
            show_label=False
        )
    
    # Footer
    gr.HTML("""
        <div class="footer-link">
            <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank">Built with anycoder</a>
        </div>
    """)
    
    # Events
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt],
        outputs=output_image,
        api_visibility="public"
    )
    
    prompt.submit(
        fn=generate_image,
        inputs=[prompt],
        outputs=output_image,
        api_visibility="public"
    )

if __name__ == "__main__":
    demo.launch(
        css=apple_css,
        footer_links=[
            {"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"}
        ]
    )