shanetx's picture
add two sections
406cd21
import os
import gradio as gr
import torch
print(f"Is CUDA available: {torch.cuda.is_available()}")
# True
print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
from diffusers import StableDiffusionImg2ImgPipeline, StableDiffusionInpaintPipeline
os.environ['GRADIO_THEME'] = 'default'
# load the pipeline
device = "cuda"
model_id_img2img = "runwayml/stable-diffusion-v1-5"
img2img_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id_img2img, torch_dtype=torch.float16)
img2img_pipe = img2img_pipe.to(device)
model_id_inpaint = "runwayml/stable-diffusion-inpainting"
inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained(model_id_inpaint, torch_dtype=torch.float16)
inpaint_pipe = inpaint_pipe.to(device)
def img2img_diff(prompt, pil_img):
img = pil_img.resize((768, 512))
return img2img_pipe(prompt=prompt, image=img, strength=0.75, guidance_scale=7.5).images[0]
def imginpaint_diff(prompt, pil_img, mask_pil_img):
return inpaint_pipe(prompt=prompt, image=pil_img, mask_image=mask_pil_img).images[0]
def header_html(title):
return f"""
<div style="text-align: center; max-width: 650px; margin: 0 auto; padding-top: 7px;">
<div
style="
display: inline-flex;
align-items: center;
gap: 0.8rem;
font-size: 1.75rem;
"
>
<h1 style="font-weight: 900; margin-bottom: 7px;">
{title}
</h1>
</div>
</div>
"""
with gr.Blocks() as block:
with gr.Group():
with gr.Box():
gr.HTML(header_html("diffusion image to image transform"))
with gr.Row():
with gr.Column():
input_img = gr.Image(type='pil', label='draft image')
with gr.Row():
input_prompt = gr.Text(lable="prompt text")
sumit_button = gr.Button("Generate image").style(
margin=False,
rounded=(False, True, True, False),
full_width=False,
)
output_img = gr.Image(type="pil")
sumit_button.click(img2img_diff, inputs=[input_prompt, input_img], outputs=[output_img])
with gr.Box():
gr.HTML(header_html("diffusion image inpaint"))
with gr.Row():
with gr.Column():
input_img = gr.Image(type='pil', label='origin image')
mask_img = gr.Image(type='pil', label='mask image')
with gr.Row():
input_prompt = gr.Text(lable="prompt text")
sumit_button = gr.Button("Generate image").style(
margin=False,
rounded=(False, True, True, False),
full_width=False,
)
output_img = gr.Image(type="pil")
sumit_button.click(imginpaint_diff, inputs=[input_prompt, input_img, mask_img], outputs=[output_img])
block.queue(concurrency_count=40, max_size=20).launch(max_threads=150)