import os
import gradio as gr
from huggingface_hub import InferenceClient
client = InferenceClient()
from gradio_imageslider import ImageSlider


def refine_image(image, prompt, negative_prompt, num_inference_steps, guidance_scale, seed, strength):
    image = client.image_to_image(
        image, 
        prompt=prompt, 
        negative_prompt=negative_prompt, 
        num_inference_steps=num_inference_steps, 
        guidance_scale=guidance_scale, 
        seed=seed, 
        model="stabilityai/stable-diffusion-xl-refiner-1.0",
        strength=strength
    )
    return image

with gr.Blocks() as demo:
    image = gr.Image(type="filepath")
    prompt = gr.Textbox(lines=3, label="Prompt")
    negative_prompt = gr.Textbox(lines=3, label="Negative Prompt")
    strength = gr.Slider(
        label="Strength",
        minimum=0,
        maximum=3,
        step=0.01,
        value=1
    )
    num_inference_steps = gr.Slider(
        label="Inference steps",
        minimum=3,
        maximum=150,
        step=1,
        value=25
    )
    guidance_scale = gr.Slider(
        label="Guidance scale",
        minimum=0.0,
        maximum=50.0,
        step=0.1,
        value=12
    )
    seed = gr.Slider(
        label="Seed",
        info="-1 denotes a random seed",
        minimum=-1,
        maximum=423538377342,
        step=1,
        value=-1
    )
    refine_btn = gr.Button("Refine")
    output = gr.Image()

    refine_btn.click(
        refine_image, 
        inputs=[image, prompt, negative_prompt, num_inference_steps, guidance_scale, seed, strength], 
        outputs=output
    )

demo.launch()