Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| from io import BytesIO | |
| import requests | |
| import PIL | |
| from PIL import Image, ImageOps | |
| import numpy as np | |
| import os | |
| import uuid | |
| import torch | |
| from torch import autocast | |
| import cv2 | |
| from matplotlib import pyplot as plt | |
| from torchvision import transforms | |
| from diffusers import DiffusionPipeline | |
| from diffusers import StableDiffusionXLInpaintPipeline | |
| api_key = os.environ.get("HF_TOKEN") | |
| # Set the model | "Stable Diffusion 1.5", "Stable Diffusion XL" | |
| pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, use_auth_token=api_key) | |
| pipe2 = StableDiffusionXLInpaintPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-0.9", torch_dtype=torch.float16, variant="fp16", use_safetensors=True, use_auth_token=api_key) | |
| #pipe.to("cuda") | |
| #pipe2.to("cuda") | |
| #from share_btn import community_icon_html, loading_icon_html, share_js | |
| def read_content(file_path: str) -> str: | |
| """read the content of target file | |
| """ | |
| with open(file_path, 'r', encoding='utf-8') as f: | |
| content = f.read() | |
| return content | |
| def load_result_to_canvas(res_to_continue): | |
| return res_to_continue | |
| def prepare_image_for_9_16(input_image_path): | |
| # Open the input image | |
| input_image = Image.open(input_image_path) | |
| # Get the width and height of the input image | |
| input_width, input_height = input_image.size | |
| # Calculate the new height for the output image | |
| output_height = int(input_width * (16 / 9)) | |
| # Calculate the vertical position for the input image in the output image | |
| y_offset = (output_height - input_height) // 2 | |
| # Create a new transparent image with the desired dimensions | |
| output_image = Image.new("RGBA", (input_width, output_height), (0, 0, 0, 0)) | |
| # Paste the input image onto the output image at the calculated position | |
| output_image.paste(input_image, (0, y_offset)) | |
| # Save the output image with transparent background | |
| output_image.save('prepared.png') | |
| return "prepared.png" | |
| def predict(img, prompt, sd_model, output_res, use_predefined_mask, inference_steps, guidance_scale, strength): | |
| # Set dimensions | "9:16 - 288x512", "9:16 - 576x1024" | |
| if output_res == "9:16 - 288x512" : | |
| w_resize = int(288) | |
| h_resize = int(512) | |
| elif output_res == "9:16 - 576x1024" : | |
| w_resize = int(576) | |
| h_resize = int(1024) | |
| print(img) | |
| init_image = Image.open(img["image"]).convert("RGB").resize((w_resize, h_resize), Image.LANCZOS) | |
| #init_image.save('init.png') | |
| #init_image = 'init.png' | |
| # Use gradio mask tool OR pre-defined mask | |
| if use_predefined_mask == False : | |
| mask = Image.open(img["mask"]).convert("RGB").resize((w_resize, h_resize), Image.LANCZOS) | |
| #mask.save('mask.png') | |
| #mask = 'mask.png' | |
| else : | |
| mask_url = "./9_16_mask.png" | |
| mask = ImageOps.invert(Image.open(mask_url).convert("RGB").resize((w_resize, h_resize), Image.LANCZOS)) | |
| #mask.save('mask.png') | |
| #mask = 'mask.png' | |
| # Get result from the model | "Stable Diffusion 1.5", "Stable Diffusion XL" | |
| if sd_model == "Stable Diffusion 1.5" : | |
| pipe.to("cuda") | |
| output = pipe(prompt = prompt, | |
| image=init_image, | |
| width=w_resize, | |
| height=h_resize, | |
| mask_image=mask, | |
| num_inference_steps=inference_steps, | |
| guidance_scale=guidance_scale | |
| ) | |
| elif sd_model == "Stable Diffusion XL" : | |
| pipe2.to("cuda") | |
| output = pipe2(prompt = prompt, | |
| image=init_image, | |
| width=w_resize, | |
| height=h_resize, | |
| mask_image=mask, | |
| num_inference_steps=inference_steps, | |
| guidance_scale=guidance_scale, | |
| strength = strength | |
| ) | |
| return output.images[0] | |
| css = ''' | |
| .container {max-width: 1150px;margin: auto;padding-top: 1.5rem} | |
| #image_upload{min-height:400px} | |
| #image_upload [data-testid="image"], #image_upload [data-testid="image"] > div{min-height: 400px} | |
| #mask_radio .gr-form{background:transparent; border: none} | |
| #word_mask{margin-top: .75em !important} | |
| #word_mask textarea:disabled{opacity: 0.3} | |
| .footer {margin-bottom: 45px;margin-top: 35px;text-align: center;border-bottom: 1px solid #e5e5e5} | |
| .footer>p {font-size: .8rem; display: inline-block; padding: 0 10px;transform: translateY(10px);background: white} | |
| .dark .footer {border-color: #303030} | |
| .dark .footer>p {background: #0b0f19} | |
| .acknowledgments h4{margin: 1.25em 0 .25em 0;font-weight: bold;font-size: 115%} | |
| #image_upload .touch-none{display: flex} | |
| @keyframes spin { | |
| from { | |
| transform: rotate(0deg); | |
| } | |
| to { | |
| transform: rotate(360deg); | |
| } | |
| } | |
| #share-btn-container { | |
| display: flex; padding-left: 0.5rem !important; padding-right: 0.5rem !important; background-color: #000000; justify-content: center; align-items: center; border-radius: 9999px !important; width: 13rem; | |
| } | |
| #share-btn { | |
| all: initial; color: #ffffff;font-weight: 600; cursor:pointer; font-family: 'IBM Plex Sans', sans-serif; margin-left: 0.5rem !important; padding-top: 0.25rem !important; padding-bottom: 0.25rem !important; | |
| } | |
| #share-btn * { | |
| all: unset; | |
| } | |
| #share-btn-container div:nth-child(-n+2){ | |
| width: auto !important; | |
| min-height: 0px !important; | |
| } | |
| #share-btn-container .wrap { | |
| display: none !important; | |
| } | |
| ''' | |
| with gr.Blocks(css=css) as demo: | |
| gr.HTML(read_content("header.html")) | |
| with gr.Group(): | |
| with gr.Box(): | |
| with gr.Row(): | |
| with gr.Column(): | |
| source_img = gr.Image(label="16:9 Source", info="Please import a 16:9 aspect ratio image", source="upload", type="filepath") | |
| load_init_img_btn = gr.Button("Load init 16:9 image") | |
| prompt = gr.Textbox(placeholder = 'Your prompt (what you want in place of what is erased)', show_label=False, elem_id="input-text") | |
| sd_model = gr.Dropdown(label="Pick a SD model",choices=["Stable Diffusion 1.5", "Stable Diffusion XL"], value="Stable Diffusion 1.5") | |
| output_res = gr.Dropdown(label="Choose an output resolution",choices=["9:16 - 288x512", "9:16 - 576x1024"], value="9:16 - 288x512") | |
| use_predefined_mask = gr.Checkbox(label="Use the pre-defined mask ? ", value=False) | |
| inference_steps = gr.Slider(label="Inference steps", minimum=25, maximum=100, step=1, value=50) | |
| guidance_scale = gr.Slider(label="Guidance scale", minimum=0.0, maximum=50.0, step=0.1, value=7.5) | |
| strength = gr.Slider(label="Strength", info="Only effective with SD XL", minimum=0.00, maximum=1.00, step=0.01, value=0.80) | |
| with gr.Column(): | |
| canvas_img = gr.Image(label="Canvas",source='upload', tool='sketch', interactive=True, elem_id="output-img", type="filepath", height=770).style(height=770) | |
| inpaint_btn = gr.Button("Inpaint!").style( | |
| margin=False, | |
| rounded=(False, True, True, False), | |
| full_width=False, | |
| ) | |
| #with gr.Group(elem_id="share-btn-container", visible=False): | |
| # community_icon = gr.HTML(community_icon_html, visible=False) | |
| # loading_icon = gr.HTML(loading_icon_html, visible=False) | |
| # share_button = gr.Button("Share to community", elem_id="share-btn", visible=False) | |
| with gr.Column(): | |
| image_out = gr.Image(label="Output", type="filepath", height=770, interactive=False).style(height=770) | |
| continue_btn = gr.Button("Continue with this output") | |
| load_init_img_btn.click( | |
| fn=prepare_image_for_9_16, | |
| inputs=[source_img], | |
| outputs=[canvas_img], | |
| queue=False | |
| ) | |
| continue_btn.click( | |
| fn=load_result_to_canvas, | |
| inputs=[image_out], | |
| outputs=[canvas_img], | |
| queue=False) | |
| inpaint_btn.click( | |
| fn=predict, | |
| inputs=[canvas_img, | |
| prompt, | |
| sd_model, | |
| output_res, | |
| use_predefined_mask, | |
| inference_steps, | |
| guidance_scale, | |
| strength], | |
| outputs=[image_out] | |
| ) | |
| #share_button.click(None, [], [], _js=share_js) | |
| demo.launch() |