Update app.py

app.py

@@ -2,44 +2,63 @@ 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
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+model_id = "Disty0/LCM_SoteMix"
+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)
 
-if torch.cuda.is_available():
-    torch.cuda.max_memory_allocated(device=device)
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
-    pipe.enable_xformers_memory_efficient_attention()
-    pipe = pipe.to(device)
-else: 
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)
-    pipe = pipe.to(device)
 
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
+pipe = OVStableDiffusionPipeline.from_pretrained(model_id, compile = False, ov_config = {"CACHE_DIR":""})
 
-def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
+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=512, width=512, num_images_per_prompt=1)
+pipe.compile()
+
+
+def infer(prompt,negative_prompt):
 
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-        
-    generator = torch.Generator().manual_seed(seed)
-    
     image = pipe(
-        prompt = prompt, 
-        negative_prompt = negative_prompt,
-        guidance_scale = guidance_scale, 
-        num_inference_steps = num_inference_steps, 
-        width = width, 
-        height = height,
-        generator = generator
+        prompt = prompt+"score_8_up,score_7_up,score_6_up,score_9,score_8_up,score_7,masterpiece,best quality,source_anime,bangs,", 
+        negative_prompt = "score_6,score_5,score_4,source_furry,pathway,walkway,face mask,heterochromia,\
+            tattoos,muscular,deformed iris,deformed pupils,long body,long neck,text,error,print,signature,\
+            logo,watermark,deformed,distorted,disfigured,bad anatomy,wrong anatomy,ugly,disgusting,\
+            cropped,crooked teeth,multiple views,bad proportions,gross proportions,cloned face,\
+            worst quality,low quality,normal quality,bad quality,lowres,poorly drawn,semi-realistic,\
+            3d,render,cg,cgi,imperfect,partial,unfinished,incomplete,monochrome,grayscale,sepia,fat,\
+            wrinkle,fat leg,fat ass,blurry,hazy,sagging breasts,longbody,lowres,\
+            bad anatomy,bad hands,missing fingers,extra digit,fewer digits,worst quality,\
+            low quality,normal quality,watermark,artist name,signature,(bad anatomy)), ((bad art)),\
+            (((bad proportions))), (b&w), (black/white), (black and white), blurry, body out of frame,\
+            canvas frame, cloned face, ((close up)), cross-eye, ((deformed)), ((disfigured)), (((duplicate))), \
+            (((extra arms))), extra fingers, (((extra legs))), ((extra limbs)), (fused fingers), gross proportions, \
+            ((morbid)), (malformed limbs), ((missing arms)), ((missing legs)), mutated, mutated hands, \
+            (((mutation))), ((mutilated)), (out of frame), ((poorly drawn face)), poorly drawn feet, \
+            ((poorly drawn hands)), tiling, (too many fingers), ((ugly)), wierd colors, (((long neck))), \
+            ugly, words, wrinkles, writing",
+        width = 512, 
+        height = 512,
+        guidance_scale=1.0,
+        num_inference_steps=8,
+        num_images_per_prompt=1,
     ).images[0] 
     
     return image
 
+
 examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
+    "A cute kitten, Japanese cartoon style.",
+    "A sweet family, dad stands next to mom, mom holds baby girl.",
     "A delicious ceviche cheesecake slice",
 ]
 
@@ -50,88 +69,29 @@ css="""
 }
 """
 
-if torch.cuda.is_available():
-    power_device = "GPU"
-else:
-    power_device = "CPU"
+
+power_device = "CPU"
 
 with gr.Blocks(css=css) as demo:
     
     with gr.Column(elem_id="col-container"):
         gr.Markdown(f"""
-        # Text-to-Image Gradio Template
+        # Disty0/LCM_SoteMix 512x512
         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)
 
-        with gr.Accordion("Advanced Settings", open=False):
-            
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
-            )
-            
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
-            )
-            
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-            
-            with gr.Row():
-                
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=512,
-                )
-                
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=512,
-                )
-            
-            with gr.Row():
-                
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=0.0,
-                )
-                
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=12,
-                    step=1,
-                    value=2,
-                )
-        
         gr.Examples(
             examples = examples,
             inputs = [prompt]
@@ -139,7 +99,7 @@ with gr.Blocks(css=css) as demo:
 
     run_button.click(
         fn = infer,
-        inputs = [prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
+        inputs = [prompt],
         outputs = [result]
     )