sdart commited on
Commit
96d24c8
1 Parent(s): d7623eb
Files changed (1) hide show
  1. model.py +9 -3
model.py CHANGED
@@ -2,14 +2,20 @@ from __future__ import annotations
2
 
3
 
4
  from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
 
5
  import torch
6
  import PIL.Image
7
  import numpy as np
8
 
 
 
9
  class Model:
10
  def __init__(self):
11
  modelID = "runwayml/stable-diffusion-v1-5"
12
- pipeline = StableDiffusionPipeline.from_pretrained(modelID, torch_dtype=torch.float16)
 
 
 
13
  #self.pipe = StableDiffusionPipeline.from_pretrained(modelID)
14
  #prompt = "a photo of an astronaut riding a horse on mars"
15
  #n_prompt = "deformed, disfigured"
@@ -22,13 +28,13 @@ class Model:
22
  num_steps:int = 20,
23
  ):
24
  seed = np.random.randint(0, np.iinfo(np.int64).max)
25
- generator = torch.Generator().manual_seed(seed)
26
  return self.pipe(prompt=prompt,
27
  negative_prompt=negative_prompt,
28
  guidance_scale=guidance_scale,
29
  num_images_per_prompt=num_images,
30
  num_inference_steps=num_steps,
31
- generator=generator).images
32
 
33
 
34
 
 
2
 
3
 
4
  from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
5
+ from diffusers import DPMSolverMultistepScheduler
6
  import torch
7
  import PIL.Image
8
  import numpy as np
9
 
10
+ device = "cpu"
11
+
12
  class Model:
13
  def __init__(self):
14
  modelID = "runwayml/stable-diffusion-v1-5"
15
+ pipe = StableDiffusionPipeline.from_pretrained(modelID, torch_dtype=torch.float16)
16
+ pipe = pipe.to(device)
17
+ pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
18
+
19
  #self.pipe = StableDiffusionPipeline.from_pretrained(modelID)
20
  #prompt = "a photo of an astronaut riding a horse on mars"
21
  #n_prompt = "deformed, disfigured"
 
28
  num_steps:int = 20,
29
  ):
30
  seed = np.random.randint(0, np.iinfo(np.int64).max)
31
+ generator = torch.Generator(device).manual_seed(seed)
32
  return self.pipe(prompt=prompt,
33
  negative_prompt=negative_prompt,
34
  guidance_scale=guidance_scale,
35
  num_images_per_prompt=num_images,
36
  num_inference_steps=num_steps,
37
+ generator=generator).images
38
 
39
 
40