EulerScheduler

main.py

@@ -11,7 +11,7 @@ import json
 import uuid
 
 import torch
-from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
+from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler, EulerDiscreteScheduler
 
 
 print(f"Is CUDA available: {torch.cuda.is_available()}")
@@ -21,41 +21,20 @@ app = FastAPI()
 @app.get("/generate")
 def generate_image(prompt, inference_steps, model):
     torch.cuda.empty_cache()
-    
     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)
-    #pipe = StableDiffusionPipeline.from_pretrained(model_id)
+    pipeline = DiffusionPipeline.from_pretrained(str(model))
     
-    #pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
-    #pipe = pipe.to("cuda")
+    #pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
+    pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config)
     
-    pipeline = DiffusionPipeline.from_pretrained(str(model))
     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, num_inference_steps=int(inference_steps)).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]
-        
-
     filename = str(uuid.uuid4()) + ".jpg"
-
-    #print(f"after filename assignment")
-
     image.save(filename)
     
     print(filename)
-    #print(f"after save")
-
     
     # Data to be written
     assertion = {