updates to nvidia

Dockerfile

@@ -6,7 +6,8 @@ FROM python:3.9
 # Python and Pip
 RUN apt-get update && apt-get install -y \
     python3-pip \ 
-    wget
+    wget \
+    nvidia-container-toolkit
 
 WORKDIR /code
 

main.py

@@ -32,11 +32,11 @@ def generate_image(prompt, inference_steps, model):
     #pipe = pipe.to("cuda")
     
     pipeline = DiffusionPipeline.from_pretrained(str(model))
-    #pipeline = pipeline.to("cuda")
-    generator = torch.Generator("cpu").manual_seed(0)
+    pipeline = pipeline.to("cuda")
+    #generator = torch.Generator("gpu").manual_seed(0)
     pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
     #image = pipeline(prompt, generator=generator).images[0]
-    image = pipeline(prompt, generator=generator, num_inference_steps=int(inference_steps)).images[0]
+    image = pipeline(prompt, num_inference_steps=int(inference_steps)).images[0]
 
 
     #prompt = "a photo of an astronaut riding a horse on mars"