AlexKoff88 commited on
Commit
9691ef5
1 Parent(s): 879699a

Create SD pipeline on each request

Browse files
Files changed (1) hide show
  1. app.py +21 -24
app.py CHANGED
@@ -5,35 +5,32 @@ from diffusers import DDPMScheduler, StableDiffusionPipeline
5
 
6
  import time
7
 
8
- scheduler = DDPMScheduler(beta_start=0.00085, beta_end=0.012,
9
- beta_schedule="scaled_linear", num_train_timesteps=1000)
10
-
11
- pretrained_model_name_or_path = "svjack/Stable-Diffusion-Pokemon-en"
12
- pipe_torch_fp32 = StableDiffusionPipeline.from_pretrained(pretrained_model_name_or_path,
13
- scheduler=scheduler)
14
- pipe_torch_fp32.safety_checker = lambda images, clip_input: (images, False)
15
-
16
- pipe_fp32 = OVStableDiffusionPipeline.from_pretrained("OpenVINO/stable-diffusion-pokemons-fp32", compile=False)# scheduler=scheduler)
17
- pipe_fp32.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)
18
- pipe_fp32.compile()
19
-
20
- pipe_int8 = OVStableDiffusionPipeline.from_pretrained("OpenVINO/stable-diffusion-pokemons-quantized-aggressive", compile=False)
21
- pipe_int8.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)
22
- pipe_int8.compile()
23
-
24
- pipe_tome_int8 = OVStableDiffusionPipeline.from_pretrained("OpenVINO/stable-diffusion-pokemons-tome-quantized-aggressive", compile=False)
25
- pipe_tome_int8.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)
26
- pipe_tome_int8.compile()
27
 
28
  pipes = {
29
- "Torch fp32": pipe_torch_fp32,
30
- "OpenVINO fp32": pipe_fp32,
31
- "OpenVINO 8-bit quantized": pipe_int8,
32
- "OpenVINO merged and quantized": pipe_tome_int8
33
  }
34
 
 
 
 
 
 
35
  def generate(prompt, option, seed):
36
- pipe = pipes[option]
37
  set_seed(int(seed))
38
  start_time = time.time()
39
  if "Torch" in option:
 
5
 
6
  import time
7
 
8
+ def create_pipeline(name):
9
+ if name == "svjack/Stable-Diffusion-Pokemon-en":
10
+ scheduler = DDPMScheduler(beta_start=0.00085, beta_end=0.012,
11
+ beta_schedule="scaled_linear", num_train_timesteps=1000)
12
+ pipe = StableDiffusionPipeline.from_pretrained(name, scheduler=scheduler)
13
+ pipe.safety_checker = lambda images, clip_input: (images, False)
14
+ else:
15
+ pipe = OVStableDiffusionPipeline.from_pretrained(name, compile=False)
16
+ pipe.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)
17
+ pipe.compile()
18
+ return pipe
 
 
 
 
 
 
 
 
19
 
20
  pipes = {
21
+ "Torch fp32": "svjack/Stable-Diffusion-Pokemon-en",
22
+ "OpenVINO fp32": "OpenVINO/stable-diffusion-pokemons-fp32",
23
+ "OpenVINO 8-bit quantized": "OpenVINO/stable-diffusion-pokemons-quantized-aggressive",
24
+ "OpenVINO merged and quantized": "OpenVINO/stable-diffusion-pokemons-tome-quantized-aggressive"
25
  }
26
 
27
+ # prefetch pipelines on start
28
+ for v in pipes.values():
29
+ pipe = create_pipeline(v)
30
+ del pipe
31
+
32
  def generate(prompt, option, seed):
33
+ pipe = create_pipeline(pipes[option])
34
  set_seed(int(seed))
35
  start_time = time.time()
36
  if "Torch" in option: