ka1kuk commited on
Commit
fd60e55
1 Parent(s): 89704e9

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +2 -3
main.py CHANGED
@@ -24,9 +24,8 @@ def dummy(images, **kwargs):
24
  return images, False
25
 
26
  async def generate_image(prompt):
27
- model_id = "SG161222/Realistic_Vision_V1.4"
28
- scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
29
- pipe = await StableDiffusionPipeline.from_pretrained(model_id, scheduler=scheduler, torch_dtype=torch.float16)
30
  pipe = pipe.to("cuda")
31
  pipe.safety_checker = dummy
32
  image = await pipe(prompt).images[0]
 
24
  return images, False
25
 
26
  async def generate_image(prompt):
27
+ model_id = "runwayml/stable-diffusion-v1-5"
28
+ pipe = await StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
 
29
  pipe = pipe.to("cuda")
30
  pipe.safety_checker = dummy
31
  image = await pipe(prompt).images[0]