cuda

main.py

@@ -10,7 +10,7 @@ import os
 import json
 
 import torch
-from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
+from diffusers import DiffusionPipeline
 
 
 print(f"Is CUDA available: {torch.cuda.is_available()}")
@@ -21,18 +21,26 @@ app = FastAPI()
 def generate_image(prompt):
     print(f"Is CUDA available: {torch.cuda.is_available()}")
 
-    model_id = "CompVis/stable-diffusion-v1-4" #stabilityai/stable-diffusion-2-1
+    #model_id = "CompVis/stable-diffusion-v1-4" #stabilityai/stable-diffusion-2-1
+    model_id = "runwayml/stable-diffusion-v1-5"
 
     # Use the DPMSolverMultistepScheduler (DPM-Solver++) scheduler here instead
     #pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-    pipe = StableDiffusionPipeline.from_pretrained(model_id)
+    #pipe = StableDiffusionPipeline.from_pretrained(model_id)
+    
     #pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
-    pipe = pipe.to("cuda")
+    #pipe = pipe.to("cuda")
+    
+    pipeline = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+    pipeline = pipeline.to("cuda")
+    generator = torch.Generator("cuda").manual_seed(0)
+    image = pipeline(prompt, generator=generator).images[0]
+
 
     #prompt = "a photo of an astronaut riding a horse on mars"
 
     #image = pipe(prompt, num_inference_steps=5).images[0]
-    image = pipe(prompt).images[0]
+    #image = pipe(prompt).images[0]
         
     print(image)    
     image.save("static/ai.jpg")