jclyo1 commited on
Commit
9495165
1 Parent(s): 64e066b

removed pipe to cuda

Browse files
Files changed (1) hide show
  1. main.py +5 -2
main.py CHANGED
@@ -17,11 +17,12 @@ app = FastAPI()
17
 
18
  @app.get("/generate")
19
  def generate_image(prompt):
20
- model_id = "stabilityai/stable-diffusion-2-1" #CompVis/stable-diffusion-v1-4, runwayml/stable-diffusion-v1-5
21
 
22
  # Use the DPMSolverMultistepScheduler (DPM-Solver++) scheduler here instead
23
  pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
24
- pipe = pipe.to("cuda")
 
25
 
26
  prompt = "a photo of an astronaut riding a horse on mars"
27
 
@@ -31,6 +32,8 @@ def generate_image(prompt):
31
  image.save("static/ai.jpg")
32
  image.save("static/ai.png")
33
 
 
 
34
  @app.get("/generate-picsum")
35
  def generate_picsum(prompt):
36
  local_filename, headers = urllib.request.urlretrieve(('https://picsum.photos/id/' + prompt + '/800/800'))
 
17
 
18
  @app.get("/generate")
19
  def generate_image(prompt):
20
+ model_id = "stabilityai/stable-diffusion-2-1" #CompVis/stable-diffusion-v1-4
21
 
22
  # Use the DPMSolverMultistepScheduler (DPM-Solver++) scheduler here instead
23
  pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
24
+ #pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
25
+ #pipe = pipe.to("cuda")
26
 
27
  prompt = "a photo of an astronaut riding a horse on mars"
28
 
 
32
  image.save("static/ai.jpg")
33
  image.save("static/ai.png")
34
 
35
+
36
+
37
  @app.get("/generate-picsum")
38
  def generate_picsum(prompt):
39
  local_filename, headers = urllib.request.urlretrieve(('https://picsum.photos/id/' + prompt + '/800/800'))