app.py

import gradio as gr
import numpy as np
import random
from diffusers import DiffusionPipeline
from optimum.intel.openvino.modeling_diffusion import OVModelVaeDecoder, OVBaseModel, OVStableDiffusionPipeline
import torch
from huggingface_hub import snapshot_download
import openvino.runtime as ov
from typing import Optional, Dict



model_id = "stabilityai/sd-turbo"
#model_id = "Disty0/sotediffusion-v2" #不可

#1024*512 記憶體不足
HIGH=512
WIDTH=512

batch_size = -1

class CustomOVModelVaeDecoder(OVModelVaeDecoder):
    def __init__(
        self, model: ov.Model, parent_model: OVBaseModel, ov_config: Optional[Dict[str, str]] = None, model_dir: str = None,
    ):
        super(OVModelVaeDecoder, self).__init__(model, parent_model, ov_config, "vae_decoder", model_dir)


pipe = OVStableDiffusionPipeline.from_pretrained(
        model_id, 
        compile = False, 
        ov_config = {"CACHE_DIR":""},
        torch_dtype=torch.int8, #快
        #torch_dtype=torch.bfloat16, #中
        #variant="fp16", 
        #torch_dtype=torch.IntTensor, #慢
        use_safetensors=False,
        )

taesd_dir = snapshot_download(repo_id="deinferno/taesd-openvino")

pipe.vae_decoder = CustomOVModelVaeDecoder(model = OVBaseModel.load_model(f"{taesd_dir}/vae_decoder/openvino_model.xml"), 
                                           parent_model = pipe, 
                                           model_dir = taesd_dir
                                          )



pipe.reshape( batch_size=-1, height=HIGH, width=WIDTH, num_images_per_prompt=1)
#pipe.load_textual_inversion("./badhandv4.pt", "badhandv4")
#pipe.load_textual_inversion("./Konpeto.pt", "Konpeto")
#<shigure-ui-style>
#pipe.load_textual_inversion("sd-concepts-library/shigure-ui-style")
#pipe.load_textual_inversion("sd-concepts-library/ruan-jia")
#pipe.load_textual_inversion("sd-concepts-library/agm-style-nao")


pipe.compile()

prompt=""
negative_prompt="(worst quality, low quality, lowres, loli, kid, child), zombie, interlocked fingers, large breasts, username, watermark,"

def infer(prompt,negative_prompt):

    image = pipe(
        prompt = prompt, 
        negative_prompt = negative_prompt,
        width = WIDTH, 
        height = HIGH,
        guidance_scale=1.0,
        num_inference_steps=2,
        num_images_per_prompt=1,
    ).images[0] 
    
    return image



examples = [
    "(Digital art, highres, best quality, 8K, masterpiece, anime screencap, perfect eyes:1.4, ultra detailed:1.5),1girl,flat chest,short messy pink hair,blue eyes,tall,thick thighs,light blue hoodie,collar,light blue shirt,black sport shorts,bulge,black thigh highs,femboy,okoto no ko,smiling,blushing,looking at viewer,inside,livingroom,sitting on couch,nighttime,dark,hand_to_mouth,",
    "1girl, silver hair, symbol-shaped pupils, yellow eyes, smiling, light particles, light rays, wallpaper, star guardian, serious face, red inner hair, power aura, grandmaster1, golden and white clothes",
    "masterpiece, best quality, highres booru, 1girl, solo, depth of field, rim lighting, flowers, petals, from above, crystals, butterfly, vegetation, aura, magic, hatsune miku, blush, slight smile, close-up, against wall,",
    "((colofrul:1.7)),((best quality)), ((masterpiece)), ((ultra-detailed)), (illustration), (detailed light), (an extremely delicate and beautiful),incredibly_absurdres,(glowing),(1girl:1.7),solo,a beautiful girl,(((cowboy shot))),standding,((Hosiery)),((beautiful off-shoulder lace-trimmed layered strapless dress+white stocking):1.25),((Belts)),(leg loops),((Hosiery)),((flower headdress)),((long white hair)),(((beautiful eyes))),BREAK,((english text)),(flower:1.35),(garden),(((border:1.75))),",
]

css="""
#col-container {
    margin: 0 auto;
    max-width: 520px;
}
"""


power_device = "CPU"

with gr.Blocks(css=css) as demo:
    
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""
        # Disty0/LCM_SoteMix {WIDTH}x{HIGH}
        Currently running on {power_device}.
        """)
        
        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )         
            run_button = gr.Button("Run", scale=0)
        
        result = gr.Image(label="Result", show_label=False)

        gr.Examples(
            examples = examples,
            fn = infer,
            inputs = [prompt],
            outputs = [result]
        )

    run_button.click(
        fn = infer,
        inputs = [prompt],
        outputs = [result]
    )

demo.queue().launch()