Spaces:
Runtime error
Runtime error
File size: 1,764 Bytes
3d5e9d1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | 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()
|