| 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_image, surgery_option, additional_info): | |
| if uploaded_image is not None: | |
| pil_image = uploaded_image.convert("RGB") | |
| 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 = ["Facelift", "Rhinoplasty", "Chin Augmentation"] | |
| iface = gr.Interface( | |
| fn=generate_image, | |
| inputs=[ | |
| gr.Image(type="pil", label="Upload an Image"), | |
| gr.Dropdown(surgery_options, label="Select Surgery Option"), | |
| gr.Textbox(label="Additional Information") | |
| ], | |
| outputs=gr.Image(type="pil", label="Generated Image"), | |
| title="AI.HAIR - Version of Plastic Surgery by Bulut", | |
| description="Visualize the results of a surgical procedure before it happens." | |
| ) | |
| iface.launch() | |