flask serve

serve.py

@@ -0,0 +1,176 @@
+from compel import Compel, ReturnedEmbeddingsType
+import logging
+from abc import ABC
+
+import diffusers
+import torch
+from diffusers import StableDiffusionXLPipeline
+
+import numpy as np
+import threading
+
+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
+import sys
+from flask import Flask, request, jsonify
+
+logger = logging.getLogger(__name__)
+logger.info("Diffusers version %s", diffusers.__version__)
+
+class DiffusersHandler(ABC):
+    """
+    Diffusers handler class for text to image generation.
+    """
+
+    def __init__(self):
+        self.initialized = False
+
+    def initialize(self, properties):
+        """In this initialize function, the Stable Diffusion model is loaded and
+        initialized here.
+        Args:
+            ctx (context): It is a JSON Object containing information
+            pertaining to the model artefacts parameters.
+        """
+        
+        logger.info("Loading diffusion model")
+        logger.info("I'm totally new and updated")
+
+
+        device_str = "cuda:" + str(properties.get("gpu_id")) if torch.cuda.is_available() and properties.get("gpu_id") is not None else "cpu"
+        
+        print("my device is " + device_str)
+        self.device = torch.device(device_str)
+        self.pipe = StableDiffusionXLPipeline.from_pretrained(
+            sys.argv[1],
+            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")
+
+        self.initialized = True
+
+    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
+            function.
+        Returns:
+            list : The preprocess function returns a list of prompts.
+        """
+        logger.info("Received requests: '%s'", raw_requests)
+        self.working = True
+                
+        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, request):
+        """Generates the image relevant to the received text.
+        Args:
+            inputs (list): List of Text from the pre-process function is passed here
+        Returns:
+            list : It returns a list of the generate images for the input text
+        """
+
+        # Handling inference for sequence_classification.
+        compel = Compel(tokenizer=[self.pipe.tokenizer, self.pipe.tokenizer_2] , text_encoder=[self.pipe.text_encoder, self.pipe.text_encoder_2], returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED, requires_pooled=[False, True])
+        
+        self.prompt = request.pop("prompt")
+        conditioning, pooled = compel(self.prompt)
+
+        # Handling inference for sequence_classification.
+        inferences = self.pipe(
+            prompt_embeds=conditioning,
+            pooled_prompt_embeds=pooled,
+            **request
+        ).images
+        
+        logger.info("Generated image: '%s'", inferences)
+        return inferences
+
+    def postprocess(self, inference_outputs):
+        """Post Process Function converts the generated image into Torchserve readable format.
+        Args:
+            inference_outputs (list): It contains the generated image of the input text.
+        Returns:
+            (list): Returns a list of the images.
+        """
+        bucket_name = "outputs-storage-prod"
+        client = storage.Client()
+        self.working = False
+        bucket = client.get_bucket(bucket_name)
+        outputs = []
+        for image in inference_outputs:
+            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
+
+
+app = Flask(__name__)
+
+# Initialize the handler on startup
+gpu_count = torch.cuda.device_count()
+if gpu_count == 0:
+    raise ValueError("No GPUs available!")
+
+handlers = [DiffusersHandler() for i in range(gpu_count)]
+for i in range(gpu_count):
+    handlers[i].initialize({"gpu_id": i})
+
+handler_lock = threading.Lock()
+handler_index = 0
+
+@app.route('/generate', methods=['POST'])
+def generate_image():
+    global handler_index
+    try:
+        # Extract raw requests from HTTP POST body
+        raw_requests = request.json
+
+        with handler_lock:
+            selected_handler = handlers[handler_index]
+            handler_index = (handler_index + 1) % gpu_count  # Rotate to the next handler
+
+        processed_request = selected_handler.preprocess([raw_requests])
+        inferences = selected_handler.inference(processed_request)
+        outputs = selected_handler.postprocess(inferences)
+
+        return jsonify({"image_urls": outputs})
+    except Exception as e:
+        logger.error("Error during image generation: %s", str(e))
+        return jsonify({"error": "Failed to generate image", "details": str(e)}), 500
+
+if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=3000, threaded=True)