import gradio as gr
from PIL import Image
import torch
from diffusers import DiffusionPipeline
from diffusers.utils import load_image
import io

pipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0")

def generate_image(uploaded_file, surgery_option, additional_info):
    if uploaded_file is not None:
        pil_image = Image.open(io.BytesIO(uploaded_file.read()))
        init_image = load_image(pil_image).convert("RGB")
        prompt = f"generate image of how this person would look after {surgery_option} also use this additional information {additional_info}"
        image = pipeline(prompt, image=init_image).images[0]
        return image
    else:
        return "No image uploaded."

surgery_options = ["Face Lift", "Nose Correction", "Chin Implant"]

iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.inputs.Image(type="file", label="Upload an Image"),
        gr.inputs.Dropdown(surgery_options, label="Select Surgery Option"),
        gr.inputs.Textbox(label="Additional Information")
    ],
    outputs=gr.outputs.Image(type="pil", label="Generated Image"),
    title="SurgiLook.ai",
    description="Visualize the results of a surgical procedure before it happens."
)

iface.launch()