app.py

import os
import numpy as np
import base64
import io
import requests
from io import BytesIO
os.system("pip install gradio==4.37.2")
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)

clear_flag = False

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, input_type):
    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, input_type)

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, input_type):
    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, input_type)

def get_dark_mode():
    return """
    () => {
        document.body.classList.toggle('dark');
    }
    """
    
def pil_image_to_data_uri(img, format="PNG"):
        buffered = BytesIO()
        img.save(buffered, format=format)
        img_str = base64.b64encode(buffered.getvalue()).decode()
        return f"data:image/{format.lower()};base64,{img_str}"

def clear_session(image):
    global clear_flag 
    clear_flag = True
    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, input_type):
    global clear_flag
    if clear_flag:
        clear_flag = False  # Reset the flag after handling the clear action
        return gr.update(value=None)
    else:
        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, input_type)
        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, input_type)

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):
    global clear_flag
    if clear_flag:
        clear_flag = False  # Reset the flag after handling the clear action
        return gr.update(value=None)  # Ensure nothing is processed if clear flag is true
    else:
        if input_img is None:
            return gr.update(value=None)
        
        img = Image.open(BytesIO(requests.get(input_img).content)).convert('RGB')

        # Read the image using Pillow
        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, 'URL'


with gr.Blocks(css=css, theme="NoCrypt/miku@1.2.1") as demo:
    gr.HTML(
                """
                <!DOCTYPE html>
                <html lang="en">
                <head>
                    <meta charset="UTF-8">
                    <meta name="viewport" content="width=device-width, initial-scale=1.0">
                    <title>S2I-Artwork Animation</title>
                    <style>
                
                        @keyframes blinkCursor {
                            from { border-right-color: rgba(255, 255, 255, 0.75); }
                            to { border-right-color: transparent; }
                        }
                
                        
                
                        @keyframes fadeIn {
                            0% { opacity: 0; transform: translateY(-10px); }
                            100% { opacity: 1; transform: translateY(0); }
                        }
                
                        @keyframes bounce {
                            0%, 20%, 50%, 80%, 100% {
                                transform: translateY(0);
                            }
                            40% {
                                transform: translateY(-10px);
                            }
                            60% {
                                transform: translateY(-5px);
                            }
                        }
                        .typewriter h1 {
                            overflow: hidden;
                            border-right: .15em solid rgba(255, 255, 255, 0.75);
                            white-space: nowrap;
                            margin: 0 auto;
                            letter-spacing: .15em;
                            animation: 
                                zoomInOut 4s infinite;
                        }
                        .animated-heading {
                            animation: fadeIn 2s ease-in-out;
                        }
                
                        .animated-link {
                            display: inline-block;
                            animation: bounce 3s infinite;
                        }
                    </style>
                </head>
                <body>
                    <div>
                        <div class="typewriter">
                            <h1 style="display: flex; align-items: center; justify-content: center; margin-bottom: 10px; text-align: center;">
                                <img src="https://imgur.com/H2SLps2.png" alt="icon" style="margin-left: 10px; height: 30px;">
                                S2I-Artwork 
                                <img src="https://imgur.com/cNMKSAy.png" alt="icon" style="margin-left: 10px; height: 30px;">: 
                                Personalized Sketch-to-Art 🧨 Diffusion Models 
                                <img src="https://imgur.com/yDnDd1p.png" alt="icon" style="margin-left: 10px; height: 30px;">
                            </h1>
                        </div>
                        <h3 class="animated-heading" style="text-align: center; margin-bottom: 10px;">Authors: Vo Nguyen An Tin, Nguyen Thiet Su</h3>
                        <h4 class="animated-heading" style="margin-bottom: 10px;">*This project is the fine-tuning task with LorA on large datasets included: COCO-2017, LHQ, Danbooru, LandScape and Mid-Journey V6</h4>
                        <h4 class="animated-heading" style="margin-bottom: 10px;">* We public 2 sketch2image-models-lora training on 30K and 60K steps with skip-connection and Transformers Super-Resolution variables</h4>
                        <h4 class="animated-heading" style="margin-bottom: 10px;">* 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</h4>
                        <h4 class="animated-heading" style="margin-bottom: 10px;">* View the full code project: 
                            <a class="animated-link" href="https://github.com/aihacker111/S2I-Artwork-Sketch-to-Image/" target="_blank">GitHub Repository</a>
                        </h4>
                        <h4 class="animated-heading" style="margin-bottom: 10px;">
                            <a class="animated-link" href="https://github.com/aihacker111/S2I-Artwork-Sketch-to-Image/" target="_blank">
                                <img src="https://cdn.buymeacoffee.com/buttons/default-orange.png" alt="Buy Me A Coffee" height="41" width="100">
                            </a>
                        </h4>
                    </div>
                </body>
                </html>
                """
            )
    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
            )
            with gr.Group():
                with gr.Row():
                    url_image = gr.Textbox(label="Image URLS", value="")
                    type_image = gr.Radio(
                                    choices=["RGB", "SKETCH"],
                                    value="SKETCH",
                                    label="Type of Image (Color Image or Sketch Image)",
                                    interactive=True)
                    with gr.Row():
                            ui_mode = gr.Radio(
                                    choices=["Light Mode", "Dark Mode"],
                                    value="Light Mode",
                                    label="Switch Light/Dark Mode UI",
                                    interactive=True)
                            zero_gpu_options = gr.Radio(
                                        choices=["GPU", "CPU"],
                                        value="GPU",
                                        label="GPU & CPU Options Spaces",
                                        interactive=True)
                            model_options = gr.Radio(
                                        choices=["350k", "350k-adapter"],
                                        value="350k-adapter",
                                        label="Type Sketch2Image models",
                                        interactive=True)
                            half_model = gr.Radio(
                                        choices=["float32", "float16"],
                                        value="float16",
                                        label="Demo Speed",
                                        interactive=True)

        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.Group():
                prompt = gr.Textbox(label="Personalized Text", value="", show_label=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')
                with gr.Accordion("S2I Advances Option", open=True):
                        with gr.Row():
                                # input_type = gr.Radio(
                                #     choices=["live-sketch", "url-sketch"],
                                #     value="live-sketch",
                                #     label="Type Sketch2Image models",
                                #     interactive=True)

                                input_type = gr.Textbox(
                                    label="Check URL or Real-time Input",
                                    interactive=True)

                                prompt_quality = gr.Radio(
                                    choices=["short-sentences", "long-sentences"],
                                    value="short-sentences",
                                    label="Long/Short of Text Prompt",
                                    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)

                                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, prompt_quality]
    outputs = [result]
    prompt.submit(fn=assign_gpu, inputs=inputs, outputs=outputs, api_name=False)


    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=[image], outputs=[image, result])
    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)
    url_image.submit(fn=get_meta_from_image, inputs=[url_image, type_image], outputs=[image, input_type])
    url_image.change(fn=get_meta_from_image, inputs=[url_image, type_image], outputs=[image, input_type])
if __name__ == '__main__':
    demo.queue()
    demo.launch(debug=True, share=False)