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,
    negative_prompt: str,
    num_inference_steps: int,
    true_cfg_scale: float,
    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=negative_prompt if negative_prompt.strip() else "",
        num_inference_steps=int(num_inference_steps),
        true_cfg_scale=true_cfg_scale
    ).images[0]
    
    return image

# Apple-inspired CSS
apple_css = """
/* Apple-style design system */
@import url('https://fonts.googleapis.com/css2?family=SF+Pro+Display:wght@300;400;500;600;700&family=SF+Pro+Text:wght@400;500;600&display=swap');

:root {
    --apple-blue: #0071e3;
    --apple-blue-hover: #0077ED;
    --apple-gray-1: #f5f5f7;
    --apple-gray-2: #e8e8ed;
    --apple-gray-3: #d2d2d7;
    --apple-gray-4: #86868b;
    --apple-gray-5: #1d1d1f;
    --apple-white: #ffffff;
    --apple-black: #000000;
    --apple-shadow: 0 4px 16px rgba(0, 0, 0, 0.08);
    --apple-shadow-hover: 0 8px 32px rgba(0, 0, 0, 0.12);
    --apple-radius: 12px;
    --apple-radius-lg: 18px;
    --apple-transition: all 0.3s cubic-bezier(0.25, 0.1, 0.25, 1);
}

/* Base styles */
.gradio-container {
    font-family: -apple-system, BlinkMacSystemFont, 'SF Pro Display', 'SF Pro Text', 'Helvetica Neue', Helvetica, Arial, sans-serif !important;
    background: linear-gradient(180deg, var(--apple-gray-1) 0%, var(--apple-white) 100%) !important;
    min-height: 100vh;
    max-width: 1200px !important;
    margin: 0 auto !important;
    padding: 40px 20px !important;
}

/* Header styling */
.header-container {
    text-align: center;
    padding: 60px 20px 40px;
    margin-bottom: 40px;
}

.header-container h1 {
    font-size: 56px !important;
    font-weight: 600 !important;
    letter-spacing: -0.02em !important;
    color: var(--apple-gray-5) !important;
    margin-bottom: 16px !important;
    line-height: 1.1 !important;
}

.header-container p {
    font-size: 21px !important;
    font-weight: 400 !important;
    color: var(--apple-gray-4) !important;
    max-width: 600px;
    margin: 0 auto !important;
    line-height: 1.5 !important;
}

.header-container a {
    color: var(--apple-blue) !important;
    text-decoration: none !important;
    transition: var(--apple-transition);
}

.header-container a:hover {
    text-decoration: underline !important;
}

/* Card-like containers */
.input-section, .output-section {
    background: var(--apple-white) !important;
    border-radius: var(--apple-radius-lg) !important;
    padding: 32px !important;
    box-shadow: var(--apple-shadow) !important;
    border: 1px solid var(--apple-gray-2) !important;
    transition: var(--apple-transition);
}

.input-section:hover, .output-section:hover {
    box-shadow: var(--apple-shadow-hover) !important;
}

/* Textbox styling */
.gradio-textbox textarea, .gradio-textbox input {
    font-family: inherit !important;
    font-size: 17px !important;
    border: 2px solid var(--apple-gray-2) !important;
    border-radius: var(--apple-radius) !important;
    padding: 16px !important;
    background: var(--apple-gray-1) !important;
    transition: var(--apple-transition);
    color: var(--apple-gray-5) !important;
}

.gradio-textbox textarea:focus, .gradio-textbox input:focus {
    border-color: var(--apple-blue) !important;
    background: var(--apple-white) !important;
    outline: none !important;
    box-shadow: 0 0 0 4px rgba(0, 113, 227, 0.15) !important;
}

.gradio-textbox label {
    font-size: 14px !important;
    font-weight: 600 !important;
    color: var(--apple-gray-5) !important;
    letter-spacing: -0.01em !important;
    margin-bottom: 8px !important;
}

/* Slider styling */
.gradio-slider input[type="range"] {
    height: 6px !important;
    background: var(--apple-gray-2) !important;
    border-radius: 3px !important;
}

.gradio-slider input[type="range"]::-webkit-slider-thumb {
    width: 22px !important;
    height: 22px !important;
    background: var(--apple-white) !important;
    border: 2px solid var(--apple-gray-3) !important;
    border-radius: 50% !important;
    box-shadow: 0 2px 8px rgba(0, 0, 0, 0.15) !important;
    cursor: pointer !important;
    transition: var(--apple-transition);
}

.gradio-slider input[type="range"]::-webkit-slider-thumb:hover {
    transform: scale(1.1);
    border-color: var(--apple-blue) !important;
}

.gradio-slider label {
    font-size: 14px !important;
    font-weight: 600 !important;
    color: var(--apple-gray-5) !important;
}

.gradio-slider .wrap {
    background: transparent !important;
}

/* Button styling */
.gradio-button {
    font-family: inherit !important;
    font-size: 17px !important;
    font-weight: 500 !important;
    border-radius: 980px !important;
    padding: 16px 32px !important;
    transition: var(--apple-transition);
    border: none !important;
    cursor: pointer !important;
    letter-spacing: -0.01em !important;
}

.gradio-button.primary {
    background: var(--apple-blue) !important;
    color: var(--apple-white) !important;
}

.gradio-button.primary:hover {
    background: var(--apple-blue-hover) !important;
    transform: scale(1.02);
}

.gradio-button.primary:active {
    transform: scale(0.98);
}

.gradio-button.secondary {
    background: var(--apple-gray-1) !important;
    color: var(--apple-blue) !important;
}

.gradio-button.secondary:hover {
    background: var(--apple-gray-2) !important;
}

/* Image output styling */
.gradio-image {
    border-radius: var(--apple-radius-lg) !important;
    overflow: hidden !important;
    background: var(--apple-gray-1) !important;
    border: 1px solid var(--apple-gray-2) !important;
}

.gradio-image img {
    border-radius: var(--apple-radius) !important;
}

/* Row and column spacing */
.gradio-row {
    gap: 16px !important;
}

.gradio-column {
    gap: 16px !important;
}

/* Section labels */
.section-label {
    font-size: 13px !important;
    font-weight: 600 !important;
    color: var(--apple-gray-4) !important;
    text-transform: uppercase !important;
    letter-spacing: 0.02em !important;
    margin-bottom: 16px !important;
}

/* Slider container */
.slider-row {
    background: var(--apple-gray-1) !important;
    border-radius: var(--apple-radius) !important;
    padding: 20px !important;
    margin-top: 8px !important;
}

/* Footer link styling */
.built-with {
    text-align: center;
    padding: 40px 20px;
    color: var(--apple-gray-4);
    font-size: 14px;
}

.built-with a {
    color: var(--apple-blue) !important;
    text-decoration: none !important;
}

.built-with a:hover {
    text-decoration: underline !important;
}

/* Progress bar */
.progress-bar {
    background: var(--apple-gray-2) !important;
    border-radius: 4px !important;
}

.progress-bar > div {
    background: linear-gradient(90deg, var(--apple-blue), #34C759) !important;
    border-radius: 4px !important;
}

/* Placeholder text */
::placeholder {
    color: var(--apple-gray-4) !important;
    opacity: 1 !important;
}

/* Remove default borders and shadows */
.block {
    border: none !important;
    box-shadow: none !important;
}

/* Group styling */
.group {
    background: var(--apple-white) !important;
    border-radius: var(--apple-radius-lg) !important;
    padding: 24px !important;
    box-shadow: var(--apple-shadow) !important;
}

/* Accordion styling */
.gradio-accordion {
    border: 1px solid var(--apple-gray-2) !important;
    border-radius: var(--apple-radius) !important;
    overflow: hidden !important;
}

.gradio-accordion > button {
    background: var(--apple-gray-1) !important;
    font-weight: 500 !important;
}

/* Hide default footer items we don't want */
footer {
    opacity: 0.7;
}
"""

with gr.Blocks(
    title="OVIS Image Generator"
) as demo:
    # Header
    gr.HTML("""
        <div class="header-container">
            <h1>OVIS Image</h1>
            <p>Create stunning images with OVIS-Image-7B. Simply describe what you imagine.</p>
            <p style="margin-top: 12px; font-size: 14px;">
                <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank">Built with anycoder</a>
            </p>
        </div>
    """)
    
    with gr.Row():
        # Left column - Input
        with gr.Column(scale=1, elem_classes="input-section"):
            gr.HTML('<div class="section-label">Create</div>')
            
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Describe the image you want to create...",
                lines=4,
                elem_id="prompt-input"
            )
            
            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                placeholder="What to avoid (optional)...",
                lines=2
            )
            
            gr.HTML('<div class="section-label" style="margin-top: 24px;">Settings</div>')
            
            with gr.Group(elem_classes="slider-row"):
                num_steps = gr.Slider(
                    label="Inference Steps",
                    minimum=10,
                    maximum=100,
                    value=50,
                    step=5,
                    info="More steps = higher quality, slower generation"
                )
                
                cfg_scale = gr.Slider(
                    label="CFG Scale",
                    minimum=1.0,
                    maximum=15.0,
                    value=5.0,
                    step=0.5,
                    info="How closely to follow your prompt"
                )
            
            generate_btn = gr.Button(
                "Generate Image",
                variant="primary",
                size="lg",
                elem_id="generate-btn"
            )
        
        # Right column - Output
        with gr.Column(scale=1, elem_classes="output-section"):
            gr.HTML('<div class="section-label">Result</div>')
            
            output_image = gr.Image(
                label="Generated Image",
                type="pil",
                elem_id="output-image",
                height=512
            )
    
    # Event handlers
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, num_steps, cfg_scale],
        outputs=output_image,
        api_visibility="public"
    )
    
    prompt.submit(
        fn=generate_image,
        inputs=[prompt, negative_prompt, num_steps, cfg_scale],
        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"}
        ]
    )