last shot

Dockerfile

@@ -1,7 +1,7 @@
+FROM nvidia/cuda:11.0-cudnn8-runtime-ubuntu18.04
 #FROM python:3.9
 
 # Using the Ubuntu image (our OS)
-FROM ubuntu:latest
 # Update package manager (apt-get) 
 # and install (with the yes flag `-y`)
 # Python and Pip

main.py

@@ -23,7 +23,7 @@ def generate_image(prompt):
     #pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
     pipe = StableDiffusionPipeline.from_pretrained(model_id)
     #pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
-    #pipe = pipe.to("cuda")
+    pipe = pipe.to("cuda")
 
     #prompt = "a photo of an astronaut riding a horse on mars"