jclyo1 commited on
Commit
94af471
1 Parent(s): 7005662
Files changed (1) hide show
  1. main.py +1 -1
main.py CHANGED
@@ -20,7 +20,7 @@ def generate_image(prompt, inference_steps, model):
20
  torch.cuda.empty_cache()
21
  print(f"Is CUDA available: {torch.cuda.is_available()}")
22
 
23
- pipeline = StableDiffusionPipeline.from_pretrained(str(model), torch_dtype=torch.float16)
24
 
25
  #pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
26
  #another comment
 
20
  torch.cuda.empty_cache()
21
  print(f"Is CUDA available: {torch.cuda.is_available()}")
22
 
23
+ pipeline = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16)
24
 
25
  #pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
26
  #another comment