gcs upload in handler

stable_diffusion_handler.py

@@ -8,6 +8,13 @@ from diffusers import StableDiffusionXLPipeline
 from ts.torch_handler.base_handler import BaseHandler
 import numpy as np
 
+import base64
+from io import BytesIO
+from PIL import Image
+import numpy as np
+import uuid
+from tempfile import TemporaryFile
+from google.cloud import storage
 
 logger = logging.getLogger(__name__)
 logger.info("Diffusers version %s", diffusers.__version__)
@@ -37,17 +44,21 @@ class DiffusersHandler(BaseHandler, ABC):
         device_str = "cuda:" + str(properties.get("gpu_id")) if torch.cuda.is_available() and properties.get("gpu_id") is not None else "cpu"
         
         self.device = torch.device(device_str)
-        self.pipe = StableDiffusionXLPipeline.from_pretrained("./")
+        self.pipe = StableDiffusionXLPipeline.from_pretrained(
+            "./",
+            torch_dtype=torch.float16,
+            use_safetensors=True,
+        )
         
         logger.info("moving model to device: %s", device_str)
         self.pipe.to(self.device)
-        
+                
         logger.info(self.device)
         logger.info("Diffusion model from path %s loaded successfully", model_dir)
 
         self.initialized = True
 
-    def preprocess(self, requests):
+    def preprocess(self, raw_requests):
         """Basic text preprocessing, of the user's prompt.
         Args:
             requests (str): The Input data in the form of text is passed on to the preprocess
@@ -55,16 +66,22 @@ class DiffusersHandler(BaseHandler, ABC):
         Returns:
             list : The preprocess function returns a list of prompts.
         """
-        logger.info("Received requests: '%s'", requests)
-
-        text = requests[0]["prompt"]
-
-        logger.info("pre-processed text: '%s'", text)
-
-        return [text]
-    
+        logger.info("Received requests: '%s'", raw_requests)
+                
+        processed_request = {
+            "prompt": raw_requests[0]["prompt"],
+            "negative_prompt": raw_requests[0].get("negative_prompt"),
+            "width": raw_requests[0].get("width"),
+            "height": raw_requests[0].get("height"),
+            "num_inference_steps": raw_requests[0].get("num_inference_steps", 30),
+            "guidance_scale": raw_requests[0].get("guidance_scale", 7.5),
+        }
+        
+        logger.info("Processed request: '%s'", processed_request)
+        return processed_request
+        
 
-    def inference(self, inputs):
+    def inference(self, request):
         """Generates the image relevant to the received text.
         Args:
             inputs (list): List of Text from the pre-process function is passed here
@@ -74,7 +91,7 @@ class DiffusersHandler(BaseHandler, ABC):
 
         # Handling inference for sequence_classification.
         inferences = self.pipe(
-            inputs, guidance_scale=7.5, num_inference_steps=50
+            **request
         ).images
         
         logger.info("Generated image: '%s'", inferences)
@@ -87,7 +104,23 @@ class DiffusersHandler(BaseHandler, ABC):
         Returns:
             (list): Returns a list of the images.
         """
-        images = []
-        for image in inference_output:
-            images.append(np.array(image).tolist())
-        return images
+        bucket_name = "outputs-storage-prod"
+        client = storage.Client()
+        bucket = client.get_bucket(bucket_name)
+        outputs = []
+        for image in inference_output.images:
+            image_name = str(uuid.uuid4())
+
+            blob = bucket.blob(image_name + '.png')
+
+            with TemporaryFile() as tmp:
+                image.save(tmp, format="png")
+                tmp.seek(0)
+                blob.upload_from_file(tmp, content_type='image/png')
+
+            # generate txt file with the image name and the prompt inside
+            blob = bucket.blob(image_name + '.txt')
+            blob.upload_from_string(self.prompt)
+
+            outputs.append('https://storage.googleapis.com/' + bucket_name + '/' + image_name + '.png')
+        return outputs