gpu

main.py

@@ -19,7 +19,9 @@ app = FastAPI()
 
 @app.get("/generate")
 def generate_image(prompt):
-    model_id = "stabilityai/stable-diffusion-2-1" #CompVis/stable-diffusion-v1-4
+    print(f"Is CUDA available: {torch.cuda.is_available()}")
+
+    model_id = "CompVis/stable-diffusion-v1-4" #stabilityai/stable-diffusion-2-1
 
     # Use the DPMSolverMultistepScheduler (DPM-Solver++) scheduler here instead
     #pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
@@ -29,7 +31,8 @@ def generate_image(prompt):
 
     #prompt = "a photo of an astronaut riding a horse on mars"
 
-    image = pipe(prompt, num_inference_steps=5).images[0]
+    #image = pipe(prompt, num_inference_steps=5).images[0]
+    image = pipe(prompt).images[0]
         
     print(image)    
     image.save("static/ai.jpg")