import gradio as gr
from diffusers import AutoPipelineForText2Image, AutoencoderKL
from diffusers.utils import load_image
import torch

vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
text_pipeline = AutoPipelineForText2Image.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", vae=vae, torch_dtype=torch.float16, variant="fp16", use_safetensors=True).to("cuda")
text_pipeline.load_ip_adapter("h94/IP-Adapter", subfolder="sdxl_models", weight_name="ip-adapter_sdxl.bin")
text_pipeline.set_ip_adapter_scale(0.6)


def text_to_image(ip, prompt, neg_prompt, width, height, ip_scale, strength, guidance, steps):
    text_pipeline.set_ip_adapter_scale(ip_scale)
    
    images = text_pipeline(
        prompt=prompt,
        ip_adapter_image=ip,
        negative_prompt=neg_prompt,
        width=width,
        height=height,
        strength=strength,
        guidance_scale=guidance,
        num_inference_steps=steps,
    ).images
    
    return images[0]


with gr.Blocks() as demo:
    gr.Markdown("""
    # IP-Adapter Playground

    by [Tony Assi](https://www.tonyassi.com/)
    """)
    with gr.Row():
        with gr.Tab("Text-to-Image"):
            text_ip = gr.Image(label='IP-Adapter Image', type='pil')
            text_prompt = gr.Textbox(label='Prompt')
            text_button = gr.Button("Generate")
        with gr.Tab("Image-to-Image"):
            image_ip = gr.Image(label='IP-Adapter Image', type='pil')
            image_image = gr.Image(label='Image', type='pil')
            image_prompt = gr.Textbox(label='Prompt')
            image_button = gr.Button("Generate")
        with gr.Tab("Inpainting"):
            inpaint_ip = gr.Image(label='IP-Adapter Image', type='pil')
            inpaint_editor = gr.ImageEditor(label='Image + Mask')
            inpaint_prompt = gr.Textbox(label='Prompt')
            inpaint_button = gr.Button("Generate")
            
        output_image = gr.Image(label='Result')
        
    with gr.Accordion("Advanced Settings", open=False):
        neg_prompt = gr.Textbox(label='Negative Prompt', value='ugly, deformed, nsfw')
        width_slider = gr.Slider(256, 1024, value=1024, label="Width")
        height_slider = gr.Slider(256, 1024, value=1024, label="Height")
        ip_scale_slider = gr.Slider(0.0, 1.0, value=0.6, label="IP-Adapter Scale")
        strength_slider = gr.Slider(0.0, 1.0, value=0.7, label="Strength")
        guidance_slider = gr.Slider(1.0, 15.0, value=7.5, label="Guidance")
        steps_slider = gr.Slider(50, 100, value=75, label="Steps")

    text_button.click(text_to_image, inputs=[text_ip, text_prompt, neg_prompt, width_slider, height_slider, ip_scale_slider, strength_slider, guidance_slider, steps_slider], outputs=output_image)

demo.launch()