Update handler

handler.py

@@ -4,20 +4,23 @@ from typing import Dict, List, Any
 # import torch
 from datetime import datetime
 
+import torch
 
+# torch.backends.cuda.matmul.allow_tf32 = True
 
-
-import requests
-from PIL import Image
-from transformers import Blip2Processor, Blip2ForConditionalGeneration
-
+from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
 
 class EndpointHandler():
     
     def __init__(self, path=""):
-        
-        self.processor = Blip2Processor.from_pretrained(path)
-        self.model = Blip2ForConditionalGeneration.from_pretrained(path, device_map="auto")
+
+        # Use the DPMSolverMultistepScheduler (DPM-Solver++) scheduler here instead
+        self.pipe = StableDiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16)
+        self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config)
+        self.pipe = self.pipe.to("cuda")
+
+        # self.pipe.enable_attention_slicing()
+        self.pipe.enable_xformers_memory_efficient_attention()
 
         # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         # self.model.eval()
@@ -64,18 +67,16 @@ class EndpointHandler():
 
 
         # pip install accelerate
-        
-        img_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/demo.jpg' 
 
+        batch_size = data.pop("batch_size",data)
+        
         now = datetime.now()
 
-        raw_image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB')
-        
-        question = "how many dogs are in the picture?"
-        inputs = self.processor(raw_image, question, return_tensors="pt").to("cuda")
-        
-        out = self.model.generate(**inputs)
-        output_text = self.processor.decode(out[0], skip_special_tokens=True)
+        with torch.inference_mode(): 
+            prompt = "a photo of an astronaut riding a horse on mars"
+            image = pipe([prompt]*batch_size, num_inference_steps=20)
+                
+        # image.save("astronaut_rides_horse.png")
 
         current = datetime.now()
 
@@ -100,4 +101,4 @@ class EndpointHandler():
         # new_tokens = output_ids[0, len(input_ids[0]) :]
         # output_text = self.tokenizer.decode(new_tokens, skip_special_tokens=True)
         
-        return [{"gen_text":output_text, "time_elapsed": str(current-now)}]
+        return [{"batch_size":batch_size, "time_elapsed": str(current-now)}]