import gradio as gr
from PIL import Image
from models import remix_image

# Gradio Interface
with gr.Blocks(theme=gr.themes.Soft(), title="Rodin.AI Image Remixer") as demo:
    gr.HTML(
        """
        <div style="text-align: center; max-width: 700px; margin: 0 auto;">
            <h1><img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/gradio-logo.svg" alt="Gradio Logo" style="height: 1em;"> Rodin.AI Image Remixer</h1>
            <p>Upload an image and provide a text prompt to remix it using a powerful diffusion model. 
            Adjust the creativity with denoising strength and prompt adherence with guidance scale.
            </p>
            <p>Built with <a href="https://huggingface.co/spaces/akhaliq/anycoder" style="text-decoration: underline;">anycoder</a></p>
        </div>
        """
    )
    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(type="pil", label="Input Image", value="https://huggingface.co/datasets/gradio/rodin-ai/resolve/main/rodin.jpeg")
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="A high-quality photo of a medieval knight, highly detailed, realistic, cinematic lighting, dramatic",
                lines=2,
                value="A high-quality photo of a medieval knight, highly detailed, realistic, cinematic lighting, dramatic",
            )
            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                placeholder="blurry, low quality, bad anatomy, deformed, ugly",
                lines=1,
                value="blurry, low quality, bad anatomy, deformed, ugly",
            )
            guidance_scale = gr.Slider(
                minimum=1.0,
                maximum=15.0,
                value=7.0,
                step=0.5,
                label="Guidance Scale (CFG)",
                info="Higher values make the image more aligned with the prompt.",
            )
            denoising_strength = gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.7,
                step=0.05,
                label="Denoising Strength",
                info="Higher values allow more changes to the original image. Lower values keep more of the original structure.",
            )
            remix_btn = gr.Button("Remix Image", variant="primary")
        with gr.Column(scale=1):
            output_image = gr.Image(label="Remixed Image", show_share_button=True)

    remix_btn.click(
        fn=remix_image,
        inputs=[input_image, prompt, negative_prompt, guidance_scale, denoising_strength],
        outputs=output_image,
        api_name="remix_image"
    )

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