import replicate
import gradio as gr
from io import BytesIO
import base64
import os

illuse = replicate.Client(api_token=os.getenv('REPLICATE'))
model_name = "andreasjansson/illusion:75d51a73fce3c00de31ed9ab4358c73e8fc0f627dc8ce975818e653317cb919b"
example_image = "https://replicate.delivery/pbxt/hHJNV9QteKX8DK2ckkUeXsqbEIKNGFXU1fN0MJoizz3iPlOjA/output-0.png"

def generate(prompt, negative_prompt, qr_content, pattern_image, num_inference_steps, guidance_scale, width, height, seed, num_outputs, controlnet_conditioning_scale, border, qrcode_background):
    try:
        inputs = {
                'prompt': prompt,
                'negative_prompt': negative_prompt,
                'qr_code_content': qr_content,
                'num_inference_steps': num_inference_steps,
                'guidance_scale': guidance_scale,
                'width': width,
                'height': height,
                'seed': seed,
                'num_outputs': num_outputs,
                'controlnet_conditioning_scale': controlnet_conditioning_scale,
                'border': border,
                'qrcode_background': qrcode_background
            }
        if pattern_image is not None:
            inputs['image'] = open(pattern_image, 'rb')

        result = illuse.run(
            model_name,
            input=inputs
        )
        return result
    except Exception as e:
        print(e)
        gr.Error(str(e))
        return


with gr.Blocks() as demo:
    gr.Markdown("""
# Illusion Diffusion Fast demo
## powered by replicate
""")
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt")
            negative_prompt = gr.Textbox(label="Negative")
            with gr.Row():
                qr_content = gr.Textbox(label="QR Code Content", placeholder="https://youtube.com/")
                pattern_input = gr.Image(label="Pattern Image(if used QR Code Content wont be used)", type="filepath")
            with gr.Accordion("Additional Settings", open=False):
                with gr.Row():
                    num_inference_steps = gr.Slider(label="num_inference_steps", minimum=20, maximum=100, step=1, value=50)
                    guidance_scale = gr.Slider(label="guidance_scale", minimum=0.1, maximum=30, step=0.01, value=7.5)
                with gr.Row():
                    width = gr.Slider(label='width', minimum=128, maximum=1024, step=8, value=768)
                    height = gr.Slider(label='height', minimum=128, maximum=1024, step=8, value=768)
                with gr.Row():
                    seed = gr.Number(label='seed', value=-1)
                    num_outputs = gr.Slider(label="num_outputs", minimum=1, maximum=4, step=1)
                with gr.Row():
                    controlnet_conditioning_scale = gr.Slider(label="controlnet_conditioning_scale", minimum=0, maximum=4, step=1, value=1)
                    border = gr.Slider(label="border", minimum=0, maximum=4, step=1, value=4)
                qrcode_background = gr.Dropdown(label="qrcode_background", choices=['gray', 'white'], value='white')
            run_btn = gr.Button("Run", variant="primary")
        output = gr.Gallery([example_image])

    generation_event = run_btn.click(generate, inputs=[prompt, negative_prompt, qr_content, pattern_input,
                                                       num_inference_steps, guidance_scale, width, height, seed,
                                                       num_outputs, controlnet_conditioning_scale, border,
                                                       qrcode_background], outputs=output)

demo.launch(show_api=False)