Spaces:
Running
Running
File size: 2,678 Bytes
65731df |
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 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
import gradio as gr
from PIL import Image
import torch
from diffusers import StableDiffusionInpaintPipeline, StableDiffusionUpscalePipeline
import numpy as np
def process_image(image, prompt, mode, scale_factor=2):
if mode == "upscale":
# Upscale pipeline
pipeline = StableDiffusionUpscalePipeline.from_pretrained(
"stabilityai/stable-diffusion-x4-upscaler"
)
pipeline.to("cuda" if torch.cuda.is_available() else "cpu")
# Process image
upscaled_image = pipeline(
prompt=prompt,
image=image,
noise_level=20,
num_inference_steps=20
).images[0]
return upscaled_image
elif mode == "inpaint":
# Inpainting pipeline
pipeline = StableDiffusionInpaintPipeline.from_pretrained(
"runwayml/stable-diffusion-inpainting"
)
pipeline.to("cuda" if torch.cuda.is_available() else "cpu")
# Create mask for extending the image
width, height = image.size
mask = Image.new('RGB', (width, height), 'white')
# Process image
result = pipeline(
prompt=prompt,
image=image,
mask_image=mask,
num_inference_steps=20
).images[0]
return result
# Gradio Interface
def create_interface():
with gr.Blocks(title="AI Image Enhancement") as interface:
gr.Markdown("# AI Image Enhancement Studio")
gr.Markdown("Enhance, upscale, and recreate images using AI")
with gr.Row():
with gr.Column():
input_image = gr.Image(type="pil", label="Upload Image")
prompt = gr.Textbox(label="Prompt", placeholder="Describe the desired enhancement...")
mode = gr.Radio(
choices=["upscale", "inpaint"],
label="Processing Mode",
value="upscale"
)
scale_factor = gr.Slider(
minimum=2,
maximum=8,
step=2,
label="Upscale Factor",
value=2
)
process_btn = gr.Button("Process Image")
with gr.Column():
output_image = gr.Image(type="pil", label="Enhanced Result")
process_btn.click(
fn=process_image,
inputs=[input_image, prompt, mode, scale_factor],
outputs=output_image
)
return interface
if __name__ == "__main__":
interface = create_interface()
interface.launch() |