import os
import numpy as np
import io
os.system("pip install gradio==4.29.0")
os.system("pip install opencv-python")
import cv2
import gradio as gr
import random
import warnings
import spaces
from PIL import Image
from S2I import Sketch2ImageController, css, scripts
dark_mode_theme = """
function refresh() {
const url = new URL(window.location);
if (url.searchParams.get('__theme') !== 'dark') {
url.searchParams.set('__theme', 'dark');
window.location.href = url.href;
}
}
"""
os.environ["TOKENIZERS_PARALLELISM"] = "false"
warnings.filterwarnings("ignore")
controller = Sketch2ImageController(gr)
def run_gpu(options, img_init, text_init, prompt_template_init, style_name_init, seeds_init, val_r_values_init, faster_init, model_name_init, clear_flag):
return controller.artwork(options, img_init, text_init, prompt_template_init, style_name_init, seeds_init, val_r_values_init, faster_init, model_name_init, clear_flag)
def run_cpu(options, img_init, text_init, prompt_template_init, style_name_init, seeds_init, val_r_values_init, faster_init, model_name_init, clear_flag):
return controller.artwork(options, img_init, text_init, prompt_template_init, style_name_init, seeds_init, val_r_values_init, faster_init, model_name_init, clear_flag)
def get_dark_mode():
return """
() => {
document.body.classList.toggle('dark');
}
"""
def clear_session():
return gr.update(value=None), gr.update(value=None)
def assign_gpu(options, img_init, text_init, prompt_template_init, style_name_init, seeds_init, val_r_values_init, faster_init, model_name_init, clear_flag):
if options == 'GPU':
decorated_run = spaces.GPU(run_gpu)
return decorated_run(options, img_init, text_init, prompt_template_init, style_name_init, seeds_init, val_r_values_init, faster_init, model_name_init, clear_flag)
else:
return run_cpu(options, img_init, text_init, prompt_template_init, style_name_init, seeds_init, val_r_values_init, faster_init, model_name_init, clear_flag)
def read_temp_file(temp_file_wrapper):
name = temp_file_wrapper.name
with open(temp_file_wrapper.name, 'rb') as f:
# Read the content of the file
file_content = f.read()
return file_content, name
def convert_to_pencil_sketch(image):
if image is None:
raise ValueError(f"Image at path {image} could not be loaded.")
# Converting it into grayscale
gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
# Inverting the image
inverted_image = 255 - gray_image
# Blurring the image
blurred = cv2.GaussianBlur(inverted_image, (25, 25), 0)
inverted_blurred = 255 - blurred
# Creating the pencil sketch
pencil_sketch = cv2.divide(gray_image, inverted_blurred, scale=256.0)
return pencil_sketch
def get_meta_from_image(input_img, type_image):
if input_img is None:
return gr.update(value=None)
file_content, _ = read_temp_file(input_img)
# Read the image using Pillow
img = Image.open(io.BytesIO(file_content)).convert("RGB")
img_np = np.array(img)
if type_image == 'RGB':
sketch = convert_to_pencil_sketch(img_np)
processed_img = 255 - sketch
elif type_image == 'SKETCH':
processed_img = 255 - img_np
# Convert the processed image back to PIL Image
img_pil = Image.fromarray(processed_img.astype('uint8'))
return img_pil
with gr.Blocks(css=css) as demo:
gr.HTML(
"""
S2I-Artwork Animation
S2I-Artwork
:
Personalized Sketch-to-Art 🧨 Diffusion Models
Authors: Vo Nguyen An Tin, Nguyen Thiet Su
*This project is the fine-tuning task with LorA on large datasets included: COCO-2017, LHQ, Danbooru, LandScape and Mid-Journey V6
* We public 2 sketch2image-models-lora training on 30K and 60K steps with skip-connection and Transformers Super-Resolution variables
* The inference and demo time of model is faster, you can slowly in the first runtime, but after that, the time process over 1.5 ~ 2s
"""
)
with gr.Row(elem_id="main_row"):
with gr.Column(elem_id="column_input"):
gr.Markdown("## SKETCH", elem_id="input_header")
image = gr.Sketchpad(
type="pil",
height=512,
width=512,
min_width=512,
image_mode="RGBA",
show_label=False,
mirror_webcam=False,
show_download_button=True,
elem_id='input_image',
brush=gr.Brush(colors=["#000000"], color_mode="fixed", default_size=4),
canvas_size=(1024, 1024),
layers=False
)
input_image = gr.File(label='Input image')
download_sketch = gr.Button(
"Download sketch", scale=1, elem_id="download_sketch"
)
with gr.Column(elem_id="column_output"):
gr.Markdown("## IMAGE GENERATE", elem_id="output_header")
result = gr.Image(
label="Result",
height=440,
width=440,
elem_id="output_image",
show_label=False,
show_download_button=True,
)
with gr.Row():
run_button = gr.Button("Generate 🪄", min_width=5, variant='primary')
randomize_seed = gr.Button(value='\U0001F3B2', variant='primary')
clear_button = gr.Button("Reset Sketch Session", min_width=10, variant='primary')
prompt = gr.Textbox(label="Personalized Text", value="", show_label=True)
with gr.Accordion("S2I Advances Option", open=True):
with gr.Row():
ui_mode = gr.Radio(
choices=["Light Mode", "Dark Mode"],
value="Light Mode",
label="Switch Light/Dark Mode UI",
interactive=True)
type_image = gr.Radio(
choices=["RGB", "SKETCH"],
value="SKETCH",
label="Type of Image (Color Image or Sketch Image)",
interactive=True)
input_type = gr.Radio(
choices=["live-sketch", "upload"],
value="live-sketch",
label="Type Sketch2Image models",
interactive=True)
style = gr.Dropdown(
label="Style",
choices=controller.STYLE_NAMES,
value=controller.DEFAULT_STYLE_NAME,
scale=1,
)
prompt_temp = gr.Textbox(
label="Prompt Style Template",
value=controller.styles[controller.DEFAULT_STYLE_NAME],
scale=2,
max_lines=1,
)
seed = gr.Textbox(label="Seed", value='42', scale=1, min_width=50)
zero_gpu_options = gr.Radio(
choices=["GPU", "CPU"],
value="GPU",
label="GPU & CPU Options Spaces",
interactive=True)
half_model = gr.Radio(
choices=["float32", "float16"],
value="float16",
label="Demo Speed",
interactive=True)
model_options = gr.Radio(
choices=["100k", "350k"],
value="350k",
label="Type Sketch2Image models",
interactive=True)
val_r = gr.Slider(
label="Sketch guidance: ",
show_label=True,
minimum=0,
maximum=1,
value=0.4,
step=0.01,
scale=3,
)
demo.load(None, None, None, js=scripts)
ui_mode.change(None, [], [], js=get_dark_mode())
randomize_seed.click(
lambda x: random.randint(0, controller.MAX_SEED),
inputs=[],
outputs=seed,
queue=False,
api_name=False,
)
inputs = [zero_gpu_options, image, prompt, prompt_temp, style, seed, val_r, half_model, model_options, input_type]
outputs = [result, download_sketch]
prompt.submit(fn=assign_gpu, inputs=inputs, outputs=outputs, api_name=False)
input_image.change(
fn=get_meta_from_image,
inputs=[
input_image, type_image
],
outputs=[
image
]
)
style.change(
lambda x: controller.styles[x],
inputs=[style],
outputs=[prompt_temp],
queue=False,
api_name=False,
).then(
fn=assign_gpu,
inputs=inputs,
outputs=outputs,
api_name=False,
)
clear_button.click(fn=clear_session, inputs=[], outputs=[image, result]).then(
fn=assign_gpu,
inputs=inputs,
outputs=outputs,
api_name=False,
)
val_r.change(assign_gpu, inputs=inputs, outputs=outputs, queue=False, api_name=False)
run_button.click(fn=assign_gpu, inputs=inputs, outputs=outputs, api_name=False)
image.change(assign_gpu, inputs=inputs, outputs=outputs, queue=False, api_name=False)
if __name__ == '__main__':
demo.queue()
demo.launch(debug=True, share=False)