app.py

import spaces
import gradio as gr
import numpy as np
import torch

from toonmage import attention_processor as attention
from toonmage.pipeline import ToonMagePipeline
from toonmage.utils import resize_numpy_image_long, seed_everything

torch.set_grad_enabled(False)

pipeline = ToonMagePipeline()

# other params
DEFAULT_NEGATIVE_PROMPT = (
    'cross-eyed, blurry, deformed eyeballs, deformed, deformed or partially rendered eyes, partially rendered objects, low resolution, disfigured hands, ugly, mutated, glitch,'
    'watermark, text, artifacts noise, worst quality, low quality, non-HDRi, lowres, flaws, flaws in the face, flaws in the eyes, extra limbs, signature'
)


@spaces.GPU
def run(*args):
    id_image = args[0]
    supp_images = args[1]
    prompt, neg_prompt, id_scale, mode, id_mix, steps, seed, n_samples, scale, H, W = args[2:]

    pipeline.debug_img_list = []
    if mode == 'fidelity':
        attention.NUM_ZERO = 8
        attention.ORTHO = False
        attention.ORTHO_v2 = True
    elif mode == 'extremely style':
        attention.NUM_ZERO = 16
        attention.ORTHO = True
        attention.ORTHO_v2 = False
    else:
        raise ValueError

    if id_image is not None:
        id_image = resize_numpy_image_long(id_image, 1024)
        id_embeddings = pipeline.get_id_embedding(id_image)
        if supp_images is not None:
            sup_images = resize_numpy_image_long(supp_images, 1024)
            supp_id_embeddings = pipeline.get_id_embedding(supp_images)
            id_embeddings = torch.cat(
                (id_embeddings, supp_id_embeddings if id_mix else supp_id_embeddings[:, :5]), dim=1
            )
    else:
        id_embeddings = None

    seed_everything(seed)
    ims = []
    for _ in range(n_samples):
        img = pipeline.inference(prompt, (1, H, W), neg_prompt, id_embeddings, id_scale, scale, steps)[0]
        ims.append(np.array(img))

    return ims, pipeline.debug_img_list


_MARKDOWN_ = """
<h2><b>This is the demo for ToonMagev2 for Image to Image Stylization</b></h2>

**Tips**

- Some examples are provided at the botton, we recommend to try the example prompts first

- a single ID image is usually sufficient, you can also supplement with additional auxiliary images

- Two modes are offered in this space, i.e. fidelity mode and extremely style mode. In most cases, the default fidelity mode should suffice. 

  If you find that the generated results are not stylized enough, you can choose the extremely style mode.
  
  Try out with different prompts using your image and do provide your feedback.

Note: The image of the women is generated using AI, therefore it does not resemble any living person. 
  
**Demo by [Sunder Ali Khowaja](https://sander-ali.github.io) - [X](https://x.com/SunderAKhowaja) -[Github](https://github.com/sander-ali) -[Hugging Face](https://huggingface.co/SunderAli17)**
"""

theme = gr.themes.Soft(
    primary_hue="blue",
    secondary_hue="purple",
    font=[gr.themes.GoogleFont('Source Code Pro'), gr.themes.GoogleFont('Public Sans'), 'system-ui', 'sans-serif'],
)
# js_func = """
# function refresh() {
#     const url = new URL(window.location);
#     if (url.searchParams.get('__theme') !== 'dark') {
#         url.searchParams.set('__theme', 'dark');
#         window.location.href = url.href;
#     }
# }
# """


with gr.Blocks(title="ToonMagev2", theme = "Yntec/HaleyCH_Theme_Orange") as SAK:
    gr.Markdown(_MARKDOWN_)
    with gr.Row():
        with gr.Column():
            submit = gr.Button("Generate")
            with gr.Row():
                face_image = gr.Image(label="ID image (main)", sources="upload", type="numpy", height=256)
                supp_image1 = gr.Image(
                    label="Additional ID image (auxiliary)", sources="upload", type="numpy", height=256
                )
            prompt = gr.Textbox(label="Prompt", value='portrait,cinematic,wolf ears,white hair')

            neg_prompt = gr.Textbox(label="Negative Prompt", value=DEFAULT_NEGATIVE_PROMPT)
            with gr.Row():
                id_scale = gr.Slider(label="ID scale", minimum=0, maximum=5, step=0.05, value=0.8, interactive=True)
                mode = gr.Dropdown(label="mode", choices=['fidelity', 'extremely style'], value='fidelity')
                id_mix = gr.Checkbox(
                    label="ID Mix (if you want to mix two ID image, please turn this on, otherwise, turn this off)",
                    value=False,
                )
            steps = gr.Slider(label="Steps", value=4, minimum=1, maximum=8, step=1)
            seed = gr.Slider(
                label="Seed", value=42, minimum=np.iinfo(np.uint32).min, maximum=np.iinfo(np.uint32).max, step=1
            )
            n_samples = gr.Slider(label="Num samples", value=2, minimum=1, maximum=4, step=1)
            scale = gr.Slider(
                label="CFG, recommend value range [1, 1.5], 1 will be faster ",
                value=1.2,
                minimum=1,
                maximum=1.5,
                step=0.1,
            )
            with gr.Row():
                H = gr.Slider(label="Height", value=1024, minimum=512, maximum=1280, step=64)
                W = gr.Slider(label="Width", value=768, minimum=512, maximum=1280, step=64)

            gr.Markdown("## Examples")
            example_inps = [
                [
                    'portrait,cinematic,wolf ears,white hair',
                    'sample_img/sample_img_test24.jpg',
                    'fidelity',
                ]
            ]
            gr.Examples(examples=example_inps, inputs=[prompt, face_image, mode], label='realistic')

            example_inps = [
                [
                    'portrait, impressionist painting, loose brushwork, vibrant color, light and shadow play',
                    'sample_img/sample_img_test24.jpg',
                    'fidelity',
                ]
            ]
            gr.Examples(examples=example_inps, inputs=[prompt, face_image, mode], label='painting style')

            example_inps = [
                [
                    'portrait, flat papercut style, silhouette, clean cuts, paper, sharp edges, minimalist,color block,man',
                    'sample_img/sample_img_test1.jpg',
                    'fidelity',
                ]
            ]
            gr.Examples(examples=example_inps, inputs=[prompt, face_image, mode], label='papercut style')

            example_inps = [
                [
                    'woman,cartoon,solo,Popmart Blind Box, Super Mario, 3d',
                    'sample_img/sample_img_test24.jpg',
                    'fidelity',
                ]
            ]
            gr.Examples(examples=example_inps, inputs=[prompt, face_image, mode], label='3d style')
            example_inps = [
                [
                    'portrait, pixar',
                    'sample_img/sample_img_test1.jpg',
                    'fidelity',
                ]
            ]
            gr.Examples(examples=example_inps, inputs=[prompt, face_image, mode], label='portrait style')
            example_inps = [
                [
                    'a high quality digital avatar, eating icecream',
                    'sample_img/image1.png',
                    'fidelity',
                ]
            ]
            gr.Examples(examples=example_inps, inputs=[prompt, face_image, mode], label='portrait style')

            example_inps = [
                [
                    'portrait, the legend of zelda, anime',
                    'sample_img/image1.png',
                    'extremely style',
                ]
            ]
            gr.Examples(examples=example_inps, inputs=[prompt, face_image, mode], label='anime style')

            example_inps = [
                [
                    'portrait, ironman',
                    'sample_img/image1.png',
                    'sample_img/sample_img_test24.jpg',
                    'fidelity',
                    True,
                ]
            ]
            gr.Examples(examples=example_inps, inputs=[prompt, face_image, supp_image1, mode, id_mix], label='id mix')

        with gr.Column():
            output = gr.Gallery(label='Output', elem_id="gallery")
            intermediate_output = gr.Gallery(label='DebugImage', elem_id="gallery", visible=False)

    inps = [
        face_image,
        supp_image1,
        prompt,
        neg_prompt,
        id_scale,
        mode,
        id_mix,
        steps,
        seed,
        n_samples,
        scale,
        H,
        W
    ]
    submit.click(fn=run, inputs=inps, outputs=[output, intermediate_output])


SAK.launch()