import gradio as gr
from diffusers import StableDiffusionInpaintPipeline
import torch
from PIL import Image, ImageDraw
# Load a pre-trained Stable Diffusion inpainting model
pipe = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16)
pipe.to("cuda")

# Function to transform mugshot into a professional portrait
def transform_to_professional(image):
    # Define prompt for a professional portrait
    prompt = "a professional headshot of a person, studio lighting, corporate background, highly detailed, sharp focus"
    negative_prompt = "blurry, low resolution, distorted, messy background"

    # Resize and center the input image
    image = image.resize((512, 512))
    
    # Add a mask (full image editing for simplicity)
    mask = Image.new("L", image.size, 255)

    # Generate professional portrait
    result = pipe(
        prompt=prompt,
        image=image,
        mask_image=mask,
        num_inference_steps=50,
        guidance_scale=7.5,
        negative_prompt=negative_prompt,
    )

    # Return the generated image
    return result.images[0]

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# Mugshot to Professional Headshot Converter")
    gr.Markdown("Upload a mugshot, and we'll transform it into a professional headshot using Stable Diffusion.")

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Upload Mugshot")
            submit_button = gr.Button("Transform")
        with gr.Column():
            output_image = gr.Image(type="pil", label="Professional Headshot")

    submit_button.click(transform_to_professional, inputs=[input_image], outputs=[output_image])

# Launch Gradio Interface
demo.launch()