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

# Clean Apple-inspired CSS
apple_css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600&display=swap');

* {
    font-family: -apple-system, BlinkMacSystemFont, 'Inter', 'Segoe UI', Roboto, sans-serif !important;
}

.gradio-container {
    max-width: 900px !important;
    margin: 0 auto !important;
    background: #fbfbfd !important;
}

.main-header {
    text-align: center;
    padding: 48px 24px 32px;
}

.main-header h1 {
    font-size: 44px;
    font-weight: 600;
    color: #1d1d1f;
    margin: 0 0 8px 0;
    letter-spacing: -0.03em;
}

.main-header p {
    font-size: 17px;
    color: #86868b;
    margin: 0;
    font-weight: 400;
}

.main-header a {
    color: #0071e3;
    text-decoration: none;
}

.main-header a:hover {
    text-decoration: underline;
}

.card {
    background: #ffffff;
    border-radius: 16px;
    padding: 28px;
    box-shadow: 0 2px 12px rgba(0,0,0,0.04);
    border: 1px solid #e5e5e7;
}

textarea, input[type="text"] {
    background: #f5f5f7 !important;
    border: none !important;
    border-radius: 10px !important;
    padding: 14px 16px !important;
    font-size: 15px !important;
    color: #1d1d1f !important;
    transition: all 0.2s ease !important;
}

textarea:focus, input[type="text"]:focus {
    background: #ffffff !important;
    box-shadow: 0 0 0 3px rgba(0,113,227,0.2) !important;
    outline: none !important;
}

textarea::placeholder, input::placeholder {
    color: #86868b !important;
}

label span {
    font-size: 13px !important;
    font-weight: 500 !important;
    color: #1d1d1f !important;
}

button.primary {
    background: #0071e3 !important;
    color: white !important;
    border: none !important;
    border-radius: 980px !important;
    padding: 12px 28px !important;
    font-size: 15px !important;
    font-weight: 500 !important;
    cursor: pointer !important;
    transition: all 0.2s ease !important;
}

button.primary:hover {
    background: #0077ed !important;
    transform: scale(1.02);
}

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

.settings-section {
    background: #f5f5f7;
    border-radius: 12px;
    padding: 20px;
    margin-top: 16px;
}

input[type="range"] {
    accent-color: #0071e3 !important;
}

.output-image {
    border-radius: 12px;
    overflow: hidden;
    background: #f5f5f7;
    min-height: 400px;
}

.footer-text {
    text-align: center;
    padding: 24px;
    font-size: 13px;
    color: #86868b;
}

.footer-text a {
    color: #0071e3;
    text-decoration: none;
}

/* Hide unnecessary elements */
.block.padded:not(.gradio-group) {
    padding: 0 !important;
    background: transparent !important;
    border: none !important;
    box-shadow: none !important;
}

.gradio-accordion {
    border: none !important;
    background: transparent !important;
}

.gradio-accordion > .label-wrap {
    background: #f5f5f7 !important;
    border-radius: 10px !important;
    padding: 12px 16px !important;
}
"""

with gr.Blocks(title="OVIS Image Generator") as demo:
    
    # Clean header
    gr.HTML("""
        <div class="main-header">
            <h1>OVIS Image</h1>
            <p>Transform your ideas into images</p>
        </div>
    """)
    
    with gr.Column(elem_classes="card"):
        prompt = gr.Textbox(
            label="What do you want to create?",
            placeholder="A serene mountain landscape at sunset...",
            lines=3,
            show_label=True
        )
        
        with gr.Accordion("Advanced Settings", open=False):
            negative_prompt = gr.Textbox(
                label="Negative prompt",
                placeholder="Things to avoid...",
                lines=2
            )
            
            with gr.Row():
                num_steps = gr.Slider(
                    label="Steps",
                    minimum=10,
                    maximum=100,
                    value=50,
                    step=5
                )
                
                cfg_scale = gr.Slider(
                    label="CFG Scale",
                    minimum=1.0,
                    maximum=15.0,
                    value=5.0,
                    step=0.5
                )
        
        generate_btn = gr.Button(
            "Generate",
            variant="primary",
            size="lg"
        )
    
    gr.HTML("<div style='height: 16px'></div>")
    
    with gr.Column(elem_classes="card"):
        output_image = gr.Image(
            label="",
            type="pil",
            show_label=False,
            elem_classes="output-image"
        )
    
    # Footer
    gr.HTML("""
        <div class="footer-text">
            <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank">Built with anycoder</a>
        </div>
    """)
    
    # Events
    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"}
        ]
    )