AlexKoff88 commited on
Commit
2188f5e
·
1 Parent(s): b85f1a6

Added model selector

Browse files
Files changed (1) hide show
  1. app.py +23 -8
app.py CHANGED
@@ -2,14 +2,29 @@ import gradio as gr
2
  from optimum.intel.openvino import OVStableDiffusionPipeline
3
  from diffusers.training_utils import set_seed
4
 
5
- quantized_pipe = OVStableDiffusionPipeline.from_pretrained("OpenVINO/Stable-Diffusion-Pokemon-en-quantized", compile=False)
6
- quantized_pipe.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)
7
- quantized_pipe.compile()
 
 
 
 
 
 
 
 
8
 
9
  prompt = "cartoon bird"
10
 
11
- def generate(image):
12
- output = quantized_pipe(prompt, num_inference_steps=50, output_type="pil")
 
 
 
 
 
 
 
13
  return output.images[0]
14
 
15
  examples = ["cartoon bird",
@@ -18,9 +33,9 @@ examples = ["cartoon bird",
18
 
19
  gr.Interface(
20
  fn=generate,
21
- inputs=gr.inputs.Textbox(placeholder="cartoon bird",
22
- label="Prompt",
23
- lines=1),
24
  outputs=gr.outputs.Image(type="pil", label="Generated Image"),
25
  title="OpenVINO-optimized Stable Diffusion",
26
  description="This is the Optimum-based demo for optimized Stable Diffusion pipeline trained on Pokemon dataset and running with OpenVINO",
 
2
  from optimum.intel.openvino import OVStableDiffusionPipeline
3
  from diffusers.training_utils import set_seed
4
 
5
+ pipe_fp32 = OVStableDiffusionPipeline.from_pretrained("OpenVINO/stable-diffusion-pokemons-fp32", compile=False)
6
+ pipe_fp32.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)
7
+ pipe_fp32.compile()
8
+
9
+ pipe_int8 = OVStableDiffusionPipeline.from_pretrained("OpenVINO/stable-diffusion-pokemons-quantized-aggressive", compile=False)
10
+ pipe_int8.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)
11
+ pipe_int8.compile()
12
+
13
+ pipe_tome_int8 = OVStableDiffusionPipeline.from_pretrained("OpenVINO/stable-diffusion-pokemons-tome-quantized-aggressive", compile=False)
14
+ pipe_tome_int8.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)
15
+ pipe_tome_int8.compile()
16
 
17
  prompt = "cartoon bird"
18
 
19
+ pipes = {
20
+ "FP32": pipe_fp32,
21
+ "8-bit quantized": pipe_int8,
22
+ "Merged and quantized": pipe_tome_int8
23
+ }
24
+
25
+ def generate(image, option):
26
+ pipe = pipes[option]
27
+ output = pipe(prompt, num_inference_steps=50, output_type="pil")
28
  return output.images[0]
29
 
30
  examples = ["cartoon bird",
 
33
 
34
  gr.Interface(
35
  fn=generate,
36
+ inputs=[gr.inputs.Textbox(placeholder="cartoon bird", label="Prompt", lines=1),
37
+ gr.inputs.Dropdown(choices=pipes.keys(), default="Merged and quantized", label="Model version"),
38
+ ],
39
  outputs=gr.outputs.Image(type="pil", label="Generated Image"),
40
  title="OpenVINO-optimized Stable Diffusion",
41
  description="This is the Optimum-based demo for optimized Stable Diffusion pipeline trained on Pokemon dataset and running with OpenVINO",