Statically reshapes the SD modelto speed up inference

app.py

@@ -5,20 +5,21 @@ import re
 
 import torch
 from transformers import AutoModelWithLMHead, AutoTokenizer, pipeline, set_seed
+from optimum.intel.openvino import OVStableDiffusionPipeline
 
-import gradio as grad
-from diffusers import StableDiffusionPipeline
 
-tokenizer = AutoTokenizer.from_pretrained("shahp7575/gpt2-horoscopes")
-model = AutoModelWithLMHead.from_pretrained("shahp7575/gpt2-horoscopes")
+horoscope_model_id = "shahp7575/gpt2-horoscopes"
+tokenizer = AutoTokenizer.from_pretrained(horoscope_model_id)
+model = AutoModelWithLMHead.from_pretrained(horoscope_model_id)
+text_generation_pipe = pipeline("text-generation", model="Gustavosta/MagicPrompt-Stable-Diffusion", tokenizer="gpt2")
+stable_diffusion_pipe = OVStableDiffusionPipeline.from_pretrained("echarlaix/stable-diffusion-v1-5-openvino", revision="fp16", compile=False)
+height = 128
+width = 128
+stable_diffusion_pipe.reshape(batch_size=1, height=height, width=width, num_images_per_prompt=1)
+stable_diffusion_pipe.compile()
 
 def fn(sign, cat):
-    sign = "scorpio"
-
     prompt = f"<|category|> {cat} <|horoscope|> {sign}"
-
-    
-
     prompt_encoded = torch.tensor(tokenizer.encode(prompt)).unsqueeze(0)
 
     sample_outputs = model.generate(
@@ -29,18 +30,16 @@ def fn(sign, cat):
         top_p=0.95,
         temperature=0.95,
         num_beams=4,
-        num_return_sequences=4,
+        num_return_sequences=1,
     )
 
     final_out = tokenizer.decode(sample_outputs[0], skip_special_tokens=True)
     starting_text = " ".join(final_out.split(" ")[4:])
-    pipe = pipeline("text-generation", model="Gustavosta/MagicPrompt-Stable-Diffusion", tokenizer="gpt2")
-
     seed = random.randint(100, 1000000)
     set_seed(seed)
-    response = pipe(starting_text  + " " + sign + " art.", max_length=(len(starting_text) + random.randint(60, 90)), num_return_sequences=1)
-    pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
-    image = pipe(response[0]["generated_text"], num_inference_steps=5).images[0]
+    response = text_generation_pipe(starting_text + " " + sign + " art", max_length=(len(starting_text) + random.randint(60, 90)), num_return_sequences=1)
+    image = stable_diffusion_pipe(response[0]["generated_text"], height=height, width=width, num_inference_steps=30).images[0]
+
     return [image, starting_text]
 
 
@@ -52,7 +51,7 @@ with block:
             with gr.Row(elem_id="prompt-container").style(mobile_collapse=False, equal_height=True):
                 text = gr.Dropdown(
                     label="Star Sign",
-                    choices=["aries", "taurus","gemini", "cancer", "leo", "virgo", "libra", "scorpio", "sagittarius", "capricorn", "aquarius", "Pisces"],
+                    choices=["Aries", "Taurus","Gemini", "Cancer", "Leo", "Virgo", "Libra", "Scorpio", "Sagittarius", "Capricorn", "Aquarius", "Pisces"],
                     show_label=True,
                     max_lines=1,
                     placeholder="Enter your prompt",
@@ -64,7 +63,7 @@ with block:
                 )
 
                 text2 = gr.Dropdown(
-                    choices=["love", "career", "wellness"],
+                    choices=["Love", "Career", "Wellness"],
                     label="Category",
                     show_label=True,
                     max_lines=1,

requirements.txt

@@ -1,3 +1,7 @@
 transformers
 torch
-diffusers
+diffusers
+onnx
+onnxruntime
+openvino
+optimum-intel