app.py

# -*- coding: utf-8 -*-
"""Copy of Anime_Pack_Gradio.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1RxVCwOkq3Q5qlEkQxhFGeUxICBujjEjR
"""


from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-zh-en")

model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-zh-en")

import gradio as gr
import numpy as np
from PIL import Image
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, DPMSolverMultistepScheduler, StableDiffusionImg2ImgPipeline

import torch
from controlnet_aux import HEDdetector
from diffusers.utils import load_image

import concurrent.futures
from threading import Thread
from compel import Compel


from transformers import pipeline


model_ckpt = "papluca/xlm-roberta-base-language-detection"
pipe = pipeline("text-classification", model=model_ckpt)



device="cuda" if torch.cuda.is_available() else "cpu"

hidden_booster_text = ", loraeyes, beautiful face, small boobs, a cup"
hidden_negative = "big boobs, huge boobs, sexy, dirty, d cup, e cup, g cup, slutty, badhandv4, ng_deepnegative_v1_75t, worst quality, low quality, extra digits, text, signature, bad anatomy, mutated hand, error, missing finger, cropped, worse quality, bad quality, lowres, floating limbs, bad hands, anatomical nonsense"
hidden_cn_booster_text = ",loraeyes漂亮的脸,小胸,贫乳,a罩杯"
hidden_cn_negative = "大胸, ，, ！, 。, ；,巨乳,性感,脏,d罩杯,e罩杯,g罩杯,骚,骚气,badhandv4, ng_deepnegative_v1_75t"

def translate(prompt):
    trans_text = prompt
    translated = model.generate(**tokenizer(trans_text, return_tensors="pt", padding=True))
    tgt_text = [tokenizer.decode(t, skip_special_tokens=True) for t in translated]
    tgt_text = ''.join(tgt_text)[:-1]
    return tgt_text

hed = HEDdetector.from_pretrained('lllyasviel/ControlNet')

controlnet_scribble = ControlNetModel.from_pretrained(
    "lllyasviel/sd-controlnet-scribble", torch_dtype=torch.float16, safety_checker=None, requires_safety_checker=False, )

pipe_scribble = StableDiffusionControlNetPipeline.from_single_file(
    "https://huggingface.co/shellypeng/anime-god/blob/main/animeGod_v10.safetensors", controlnet=controlnet_scribble, safety_checker=None, requires_safety_checker=False,
    torch_dtype=torch.float16,
)

# pipe.load_lora_weights("shellypeng/detail-tweaker")
# pipe.load_lora_weights("shellypeng/midjourney-anime")

# pipe.load_lora_weights("shellypeng/animetarot")
# pipe.load_lora_weights("shellypeng/anime-stickers-v3")
# pipe.load_lora_weights("shellypeng/anime-magazine")

# pipe_img2img.load_lora_weights("yenojunie/slit-pupils")

# pipe_scribble.load_lora_weights("shellypeng/detail-tweaker")
# pipe_scribble.fuse_lora(lora_scale=0.1)
pipe_scribble.load_lora_weights("shellypeng/lora-eyes")
pipe_scribble.fuse_lora(lora_scale=0.1)
# pipe_scribble.load_lora_weights("shellypeng/beautiful-eyes")
# pipe_scribble.fuse_lora(lora_scale=0.1)

pipe_scribble.load_textual_inversion("shellypeng/bad-prompt")
pipe_scribble.load_textual_inversion("shellypeng/badhandv4")
# pipe.load_textual_inversion("shellypeng/easynegative")
pipe_scribble.load_textual_inversion("shellypeng/deepnegative")
pipe_scribble.load_textual_inversion("shellypeng/verybadimagenegative")
pipe_scribble.scheduler = DPMSolverMultistepScheduler.from_config(pipe_scribble.scheduler.config, use_karras_sigmas=True)
# pipe.enable_model_cpu_offload()
pipe_scribble.safety_checker = None
pipe_scribble.requires_safety_checker = False
pipe_scribble.to(device)

def scribble_to_image(text, neg_prompt_box, input_img):
    """
    pass the sd model and do scribble to image
    include Adetailer, detail tweaker lora, prompt backend include: beautiful eyes, beautiful face, beautiful hand, (maybe infer from user's prompt for gesture and facial
    expression to improve hand)
    """



# if auto detect detects chinese => auto turn on chinese prompting checkbox
    # change param "bag" below to text, image param below to input_img
    input_img = Image.fromarray(input_img)
    input_img = hed(input_img, scribble=True)
    input_img = load_image(input_img)
    # global prompt
    lang_check_label = pipe(text, top_k=1, truncation=True)[0]['label']
    lang_check_score = pipe(text, top_k=1, truncation=True)[0]['score']
    if lang_check_label == 'zh' and lang_check_score >= 0.85:
        text = translate(text)
    compel_proc = Compel(tokenizer=pipe_scribble.tokenizer, text_encoder=pipe_scribble.text_encoder)
    prompt = text + hidden_booster_text
    prompt_embeds = compel_proc(prompt)
    negative_prompt = neg_prompt_box + hidden_negative
    negative_prompt_embeds = compel_proc(negative_prompt)

    res_image0 = pipe_scribble(image=input_img, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, num_inference_steps=40).images[0]
    res_image1 = pipe_scribble(image=input_img, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, num_inference_steps=40).images[0]
    res_image2 = pipe_scribble(image=input_img, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, num_inference_steps=40).images[0]
    res_image3 = pipe_scribble(image=input_img, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, num_inference_steps=40).images[0]

    return res_image0, res_image1, res_image2, res_image3

def real_img2img_to_anime(text, neg_prompt_box, input_img):
    """
    pass the sd model and do scribble to image
    include Adetailer, detail tweaker lora, prompt backend include: beautiful eyes, beautiful face, beautiful hand, (maybe infer from user's prompt for gesture and facial
    expression to improve hand)
    """
    input_img = Image.fromarray(input_img)
    input_img = load_image(input_img)
    lang_check_label = pipe(text, top_k=1, truncation=True)[0]['label']
    lang_check_score = pipe(text, top_k=1, truncation=True)[0]['score']
    if lang_check_label == 'zh' and lang_check_score >= 0.85:
        text = translate(text)

    compel_proc = Compel(tokenizer=pipe_scribble.tokenizer, text_encoder=pipe_scribble.text_encoder)
    prompt = text + hidden_booster_text
    prompt_embeds = compel_proc(prompt)

    negative_prompt = neg_prompt_box + hidden_negative
    negative_prompt_embeds = compel_proc(negative_prompt)
    # input_img = depth_estimator(input_img)['depth']
    res_image0 = pipe_img2img(image=input_img, strength=0.6, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, num_inference_steps=40).images[0]
    res_image1 = pipe_img2img(image=input_img, strength=0.6, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, num_inference_steps=40).images[0]
    res_image2 = pipe_img2img(image=input_img, strength=0.6, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, num_inference_steps=40).images[0]
    res_image3 = pipe_img2img(image=input_img, strength=0.6, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, num_inference_steps=40).images[0]

    return res_image0, res_image1, res_image2, res_image3


theme = gr.themes.Soft(
    primary_hue="orange",
    secondary_hue="orange",
).set(
    block_background_fill='*primary_50'
)

# %cd /content/drive/MyDrive/stable-diffusion-webui-colab/stable-diffusion-webui

pipe_img2img = StableDiffusionImg2ImgPipeline.from_single_file("https://huggingface.co/shellypeng/anime-god/blob/main/animeGod_v10.safetensors",
                                                               torch_dtype=torch.float16, safety_checker=None, requires_safety_checker=False)

# pipe_img2img.load_lora_weights("shellypeng/detail-tweaker")
# pipe_img2img.fuse_lora(lora_scale=0.1)
pipe_img2img.load_lora_weights("shellypeng/lora-eyes")
pipe_img2img.fuse_lora(lora_scale=0.1)
# pipe_img2img.load_lora_weights("shellypeng/beautiful-eyes")
# pipe_img2img.fuse_lora(lora_scale=0.1)

pipe_img2img.load_textual_inversion("shellypeng/bad-prompt")
pipe_img2img.load_textual_inversion("shellypeng/badhandv4")
# pipe.load_textual_inversion("shellypeng/easynegative")
pipe_img2img.load_textual_inversion("shellypeng/deepnegative")
pipe_img2img.load_textual_inversion("shellypeng/verybadimagenegative")
pipe_img2img.scheduler = DPMSolverMultistepScheduler.from_config(pipe_img2img.scheduler.config, use_karras_sigmas=True)
# pipe.enable_model_cpu_offload()
pipe_img2img.safety_checker = None
pipe_img2img.requires_safety_checker = False
pipe_img2img.to(device)

# pipe_img2img.safety_checker = lambda images, **kwargs: (images, [False] * len(images))

# depth_estimator = pipeline('depth-estimation')

# controlnet_depth = ControlNetModel.from_pretrained(
#     "lllyasviel/sd-controlnet-depth", torch_dtype=torch.float16
# )


# # models that worked well: anime god, pastel dream,
# # https://huggingface.co/shellypeng/featureless/tree/main
# pipe_depth = StableDiffusionControlNetPipeline.from_single_file(
#     "https://huggingface.co/shellypeng/anime-god/blob/main/animeGod_v10.safetensors", controlnet=controlnet_depth,
#     torch_dtype=torch.float16,
# )
# # pipe = StableDiffusionControlNetPipeline.from_pretrained("furusu/SSD-1B-anime",
# #     torch_dtype=torch.float16
# # )

# pipe_depth.load_lora_weights("shellypeng/detail-tweaker")
# pipe_depth.fuse_lora(lora_scale=0.1)
# # pipe.load_lora_weights("shellypeng/stylized-3d")
# # pipe.load_lora_weights("shellypeng/midjourney-anime")

# # pipe.load_lora_weights("shellypeng/animetarot")
# # pipe.load_lora_weights("shellypeng/anime-stickers-v3")
# # pipe.load_lora_weights("shellypeng/anime-magazine")



# pipe_depth.load_textual_inversion("shellypeng/bad-prompt")
# pipe_depth.load_textual_inversion("shellypeng/badhandv4")
# # pipe.load_textual_inversion("shellypeng/easynegative")
# pipe_depth.load_textual_inversion("shellypeng/deepnegative")
# pipe_depth.load_textual_inversion("shellypeng/verybadimagenegative")
# pipe_depth.scheduler = DPMSolverMultistepScheduler.from_config(pipe_depth.scheduler.config, use_karras_sigmas=True)
# # pipe.enable_model_cpu_offload()
# def dummy(images, **kwargs):
#     return images, False
# pipe_depth.safety_checker = lambda images, **kwargs: (images, [False] * len(images))
# pipe_depth.to(device)
# # pipe.load_lora_weights("shellypeng/detail-tweaker", weight_name="add_detail.safetensors")

# # load textual inversion negative embeddings!!!: pipe.load_textual_inversion("sd-concepts-library/cat-toy")

# def real_to_anime(text, input_img):
#     """
#     pass the sd model and do scribble to image
#     include Adetailer, detail tweaker lora, prompt backend include: beautiful eyes, beautiful face, beautiful hand, (maybe infer from user's prompt for gesture and facial
#     expression to improve hand)
#     """
#     input_img = Image.fromarray(input_img)
#     input_img = load_image(input_img)
#     input_img = depth_estimator(input_img)['depth']
#     res_image0 = pipe_depth(prompt, input_img, negative_prompt=hidden_negative, num_inference_steps=40).images[0]
#     res_image1 = pipe_depth(prompt, input_img, negative_prompt=hidden_negative, num_inference_steps=40).images[0]
#     res_image2 = pipe_depth(prompt, input_img, negative_prompt=hidden_negative, num_inference_steps=40).images[0]
#     res_image3 = pipe_depth(prompt, input_img, negative_prompt=hidden_negative, num_inference_steps=40).images[0]

#     return res_image0, res_image1, res_image2, res_image3


# theme = gr.themes.Soft(
#     primary_hue="orange",
#     secondary_hue="orange",
# ).set(
#     block_background_fill='*primary_50'
# )

def zh_prompt_info(text, neg_text, chinese_check):
    can_raise_info = ""
    lang_check_label = pipe(text, top_k=1, truncation=True)[0]['label']
    lang_check_score = pipe(text, top_k=1, truncation=True)[0]['score']
    neg_lang_check_label = pipe(neg_text, top_k=1, truncation=True)[0]['label']
    neg_lang_check_score = pipe(neg_text, top_k=1, truncation=True)[0]['score']
    print(lang_check_label)
    if lang_check_label == 'zh' and lang_check_score >= 0.85:
        if not chinese_check:
            chinese_check = True
            can_raise_info = "zh"
        if neg_lang_check_label == 'en' and neg_lang_check_score >= 0.85:
            can_raise_info = "invalid"
            return True, can_raise_info
    elif lang_check_label == 'en' and lang_check_score >= 0.85:
        if chinese_check:
            chinese_check = False
            can_raise_info = "en"
        if neg_lang_check_label == 'zh' and neg_lang_check_score >= 0.85:
            can_raise_info = "invalid"
            return False, can_raise_info
    return chinese_check, can_raise_info
def mult_thread_img2img(prompt_box, neg_prompt_box, image_box):
    with concurrent.futures.ThreadPoolExecutor(max_workers=12000) as executor:
        future = executor.submit(real_img2img_to_anime, prompt_box, neg_prompt_box, image_box)
        image1, image2, image3, image4 = future.result()
    return image1, image2, image3, image4
def mult_thread_scribble(prompt_box, neg_prompt_box, image_box):
    with concurrent.futures.ThreadPoolExecutor(max_workers=12000) as executor:
        future = executor.submit(scribble_to_image, prompt_box, neg_prompt_box, image_box)
        image1, image2, image3, image4 = future.result()
    return image1, image2, image3, image4
def mult_thread_lang_class(prompt_box, neg_prompt_box, chinese_check):

    with concurrent.futures.ThreadPoolExecutor(max_workers=12000) as executor:
        future = executor.submit(zh_prompt_info, prompt_box, neg_prompt_box, chinese_check)
        chinese_check, can_raise_info = future.result()
    if can_raise_info == "zh":
        gr.Info("Chinese Language Detected, Switching to Chinese Prompt Mode")
    elif can_raise_info == "en":
        gr.Info("English Language Detected, Disabling Chinese Prompt Mode")
    return chinese_check

with gr.Blocks(theme=theme, css="footer {visibility: hidden}", title="ShellAI Apps") as iface:
    with gr.Tab("Animefier"):
        with gr.Row(equal_height=True):
            with gr.Column():
                with gr.Row(equal_height=True):
                    with gr.Column(scale=4):
                        prompt_box = gr.Textbox(label="Prompt", placeholder="Enter a prompt", lines=3)
                        neg_prompt_box = gr.Textbox(label="Negative Prompt", placeholder="Enter a negative prompt(things you don't want to include in the generated image)", lines=3)
                    with gr.Row(equal_height=True):
                        chinese_check = gr.Checkbox(label="Chinese Prompt Mode", info="Click here to enable Chinese Prompting（点此触发中文提示词输入）")

                image_box = gr.Image(label="Input Image", height=350)
                gen_btn = gr.Button(value="Generate")

        with gr.Row(equal_height=True):
            global image1
            global image2
            global image3
            global image4
            image1 = gr.Image(label="Result 1")
            image2 = gr.Image(label="Result 2")
            image3 = gr.Image(label="Result 3")
            image4 = gr.Image(label="Result 4")


        gr.on(triggers=[prompt_box.submit, gen_btn.click],fn=mult_thread_lang_class, inputs=[prompt_box, neg_prompt_box, chinese_check], outputs=[chinese_check], show_progress=False)
        gr.on(triggers=[prompt_box.submit, gen_btn.click],fn=mult_thread_img2img, inputs=[prompt_box, neg_prompt_box, image_box], outputs=[image1, image2, image3, image4])

    with gr.Tab("AniSketch"):
        with gr.Row(equal_height=True):
            with gr.Column():
                with gr.Row(equal_height=True):
                    with gr.Column(scale=4):
                        prompt_box = gr.Textbox(label="Prompt", placeholder="Enter a prompt", lines=3)
                        neg_prompt_box = gr.Textbox(label="Negative Prompt", placeholder="Enter a negative prompt(things you don't want to include in the generated image)", lines=3)
                    with gr.Row(equal_height=True):
                        chinese_check = gr.Checkbox(label="Chinese Prompt Mode", info="Click here to enable Chinese Prompting（点此触发中文提示词输入）")

                image_box = gr.Image(label="Input Image", height=350)
                gen_btn = gr.Button(value="Generate")
        with gr.Row(equal_height=True):
            image1 = gr.Image(label="Result 1")
            image2 = gr.Image(label="Result 2")
            image3 = gr.Image(label="Result 3")
            image4 = gr.Image(label="Result 4")

        gr.on(triggers=[prompt_box.submit, gen_btn.click],fn=mult_thread_lang_class, inputs=[prompt_box, neg_prompt_box, chinese_check], outputs=[chinese_check], show_progress=False)
        gr.on(triggers=[prompt_box.submit, gen_btn.click],fn=mult_thread_scribble, inputs=[prompt_box, neg_prompt_box, image_box], outputs=[image1, image2, image3, image4])


        # gen_btn.click(mult_thread, [prompt_box, image_box, chinese_check], [image1, image2, image3, image4, chinese_check])
iface.dependencies[0]["show_progress"] = "hidden"
iface.launch(debug=True, share=True, auth=["shenrym", "shjdqw%23-sw2&"])