loras serve

gunicorn_config.py

@@ -0,0 +1,12 @@
+# gunicorn_config.py
+import os
+
+worker_id_counter = 0
+
+def pre_fork(server, worker):
+    global worker_id_counter
+    worker_id_counter += 1
+
+def post_fork(server, worker):
+    worker_id = worker_id_counter - 1
+    os.environ['WORKER_ID'] = str(worker_id % 4)

serve_loras.py

@@ -9,6 +9,7 @@ from diffusers import StableDiffusionXLPipeline, DiffusionPipeline
 
 import numpy as np
 import threading
+import mmap
 
 import base64
 from io import BytesIO
@@ -23,6 +24,7 @@ import os
 from sequential_timer import SequentialTimer
 from safetensors.torch import load_file
 import copy
+import gc
 
 logger = logging.getLogger(__name__)
 logger.info("Diffusers version %s", diffusers.__version__)
@@ -36,23 +38,25 @@ sentry_sdk.init(
 
 LORAS_DIR = './safetensors'
 
+lora_lock = threading.Lock()
+
 handler_lock = threading.Lock()
 handler_index = 0
 
-class LoraCache():
-    def __init__(self, loras_dir: str = LORAS_DIR):
-        self.loras_dir = loras_dir
-        self.cache = {}
-
-    def load_lora(self, lora_name: str):
-        if lora_name.endswith('.safetensors'):
-            lora_name = lora_name.rstrip('.safetensors')
-        if lora_name not in self.cache:
-            lora = load_file(os.path.join(self.loras_dir, lora_name+'.safetensors'))
-            self.cache[lora_name] = lora
-        return copy.deepcopy(self.cache[lora_name])
+# class LoraCache():
+#     def __init__(self, loras_dir: str = LORAS_DIR):
+#         self.loras_dir = loras_dir
+#         self.cache = {}
+
+#     def load_lora(self, lora_name: str):
+#         if lora_name.endswith('.safetensors'):
+#             lora_name = lora_name.rstrip('.safetensors')
+#         if lora_name not in self.cache:
+#             lora = load_file(os.path.join(self.loras_dir, lora_name+'.safetensors'))
+#             self.cache[lora_name] = lora
+#         return copy.deepcopy(self.cache[lora_name])
     
-lora_cache = LoraCache()
+# lora_cache = LoraCache()
 
 class DiffusersHandler(ABC):
     """
@@ -134,7 +138,7 @@ class DiffusersHandler(ABC):
             "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", 25),
+            "num_inference_steps": raw_requests[0].get("num_inference_steps", 30),
             "guidance_scale": raw_requests[0].get("guidance_scale", 8.5)
             # "lora_weights": raw_requests[0].get("lora_name", None)
             # "cross_attention_kwargs": {"scale": raw_requests[0].get("lora_scale", 0.0)}
@@ -167,6 +171,7 @@ class DiffusersHandler(ABC):
         # 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])
         st = SequentialTimer()
         model_args, extra_args = request
+        global lora_cache
 
         use_char_lora = extra_args['char_lora'] is not None
         use_style_lora = extra_args['style_lora'] is not None
@@ -188,7 +193,8 @@ class DiffusersHandler(ABC):
         if use_style_lora:
             style_lora = os.path.join(LORAS_DIR, style_lora + '.safetensors')
             st.time("Load style lora")
-            self.pipe.load_lora_weights(style_lora)
+            with lora_lock:
+                self.pipe.load_lora_weights(style_lora)
             if use_char_lora:
                 st.time("Fuse style lora into model")
                 self.pipe.fuse_lora(lora_scale=extra_args['style_scale'], fuse_text_encoder=False)
@@ -196,7 +202,8 @@ class DiffusersHandler(ABC):
         if use_char_lora:
             char_lora = os.path.join(LORAS_DIR, char_lora + '.safetensors')
             st.time('load character lora')
-            self.pipe.load_lora_weights(char_lora)
+            with lora_lock:
+                self.pipe.load_lora_weights(char_lora)
 
         # lora_weights = model_args.pop("lora_weights")
         # if lora_weights is not None:
@@ -287,6 +294,8 @@ def generate_image():
         axiom_logger.info(message="Received request", request_id=req_id, **raw_requests)
 
         with handler_lock:
+            if handler_index == 0:
+                gc.collect()
             selected_handler = handlers[handler_index]
             handler_index = (handler_index + 1) % gpu_count  # Rotate to the next handler
             selected_handler.req_id = req_id
@@ -295,7 +304,6 @@ def generate_image():
         inferences = selected_handler.inference(processed_request)
         outputs = selected_handler.postprocess(inferences)
         selected_handler.req_id = None
-
         return jsonify({"image_urls": outputs})
     except Exception as e:
         logger.error("Error during image generation: %s", str(e))

serve_loras_prod.py

@@ -0,0 +1,318 @@
+from compel import Compel, ReturnedEmbeddingsType
+import logging
+from abc import ABC
+import uuid
+
+import diffusers
+import torch
+from diffusers import StableDiffusionXLPipeline, DiffusionPipeline
+
+import numpy as np
+import threading
+
+import base64
+from io import BytesIO
+from PIL import Image
+import numpy as np
+from tempfile import TemporaryFile
+from google.cloud import storage
+import sys
+import sentry_sdk
+from flask import Flask, request, jsonify, current_app
+import os
+from sequential_timer import SequentialTimer
+from safetensors.torch import load_file
+from dotenv import load_dotenv
+import copy
+import gc
+
+
+logger = logging.getLogger(__name__)
+logger.info("Diffusers version %s", diffusers.__version__)
+
+from axiom_logger import AxiomLogger
+axiom_logger = AxiomLogger()
+
+sentry_sdk.init(
+    dsn="https://f750d1b039d66541f344ee6151d38166@o4505891057696768.ingest.sentry.io/4506071735205888",
+)
+
+LORAS_DIR = './safetensors'
+
+load_dotenv()
+
+lora_lock = threading.Lock()
+
+# handler_lock = threading.Lock()
+# handler_index = 0
+
+# class LoraCache():
+#     def __init__(self, loras_dir: str = LORAS_DIR):
+#         self.loras_dir = loras_dir
+#         self.cache = {}
+
+#     def load_lora(self, lora_name: str):
+#         if lora_name.endswith('.safetensors'):
+#             lora_name = lora_name.rstrip('.safetensors')
+#         if lora_name not in self.cache:
+#             lora = load_file(os.path.join(self.loras_dir, lora_name+'.safetensors'))
+#             self.cache[lora_name] = lora
+#         return copy.deepcopy(self.cache[lora_name])
+    
+# lora_cache = LoraCache()
+
+class DiffusersHandler(ABC):
+    """
+    Diffusers handler class for text to image generation.
+    """
+
+    def __init__(self):
+        self.initialized = False
+        self.req_id = None
+
+    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"
+        self.device_str = device_str
+        
+        print("my device is " + device_str)
+        self.device = torch.device(device_str)
+        self.pipe = StableDiffusionXLPipeline.from_pretrained(
+            "./",
+            torch_dtype=torch.float16,
+            use_safetensors=True,
+        )
+        # self.refiner = DiffusionPipeline.from_pretrained(
+        #     "stabilityai/stable-diffusion-xl-refiner-1.0",
+        #     text_encoder_2=self.pipe.text_encoder_2,
+        #     vae=self.pipe.vae,
+        #     torch_dtype=torch.float16,
+        #     use_safetensors=True,
+        #     variant="fp16",
+        # )
+        # self.refiner.enable_model_cpu_offload(properties.get("gpu_id"))
+        # logger.info("Refiner initialized and o")
+
+        self.compel_base = 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])
+        logger.info("Compel initialized")
+        
+        # self.compel_refiner = Compel(
+        #     tokenizer=[self.refiner.tokenizer_2],
+        #     text_encoder=[self.refiner.text_encoder_2],
+        #     returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
+        #     requires_pooled=[True])
+        
+        logger.info("moving base model to device: %s", device_str)
+        self.pipe.to(self.device)
+                
+        logger.info(self.device)
+        logger.info("Diffusion model from path %s loaded successfully")
+        axiom_logger.info("Diffusion model initialized", device=self.device_str)
+
+        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
+                
+        model_args = {
+            "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", 8.5)
+            # "lora_weights": raw_requests[0].get("lora_name", None)
+            # "cross_attention_kwargs": {"scale": raw_requests[0].get("lora_scale", 0.0)}
+        }
+
+        extra_args = {
+            "seed": raw_requests[0].get("seed", None),
+            "style_lora": raw_requests[0].get("style_lora", None),
+            "style_scale": raw_requests[0].get("style_scale", 1.0),
+            "char_lora": raw_requests[0].get("char_lora", None),
+            "char_scale": raw_requests[0].get("char_scale", 1.0),
+            "scene_prompt": raw_requests[0].get("scene_prompt", None)
+        }
+
+        
+        logger.info("Processed request: '%s'", model_args)
+        axiom_logger.info("Processed request:" + str(model_args), request_id=self.req_id, device=self.device_str)
+        return model_args, extra_args
+        
+
+    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])
+        st = SequentialTimer()
+        model_args, extra_args = request
+
+        use_char_lora = extra_args['char_lora'] is not None
+        use_style_lora = extra_args['style_lora'] is not None
+
+        
+        style_lora = extra_args['style_lora']
+        char_lora = extra_args['char_lora']
+
+        cross_attention_kwargs = {"scale": extra_args['char_scale'] if use_char_lora else extra_args['style_scale']}
+
+        generator = torch.Generator(device="cuda").manual_seed(extra_args['seed']) if extra_args['seed'] else None
+
+
+        prompt = model_args.pop("prompt")
+        negative_prompt = model_args.pop('negative_prompt')
+        scene_prompt = extra_args['scene_prompt']
+        if scene_prompt:
+            prompt = f'("{prompt}", "{scene_prompt}").and()'
+        st.time("Base compel embedding")
+        conditioning, pooled = self.compel_base(prompt)
+        negative_conditioning, negative_pooled = self.compel_base(negative_prompt)
+        
+        [conditioning, negative_conditioning] = self.compel_base.pad_conditioning_tensors_to_same_length([conditioning, negative_conditioning])
+
+        if use_style_lora:
+            style_lora = os.path.join(LORAS_DIR, style_lora + '.safetensors')
+            st.time("Load style lora")
+            self.pipe.load_lora_weights(style_lora)
+            if use_char_lora:
+                st.time("Fuse style lora into model")
+                self.pipe.fuse_lora(lora_scale=extra_args['style_scale'], fuse_text_encoder=False)
+        
+        if use_char_lora:
+            char_lora = os.path.join(LORAS_DIR, char_lora + '.safetensors')
+            st.time('load character lora')
+            self.pipe.load_lora_weights(char_lora)
+
+        # lora_weights = model_args.pop("lora_weights")
+        # if lora_weights is not None:
+        #     lora_path = os.path.join(LORAS_DIR, lora_weights + '.safetensors')
+        #     logger.info('LOADING LORA FROM: ' + lora_path)
+        #     self.pipe.load_lora_weights(lora_path)
+
+        # Handling inference for sequence_classification.
+        st.time("base model inference")
+        inferences = self.pipe(
+            prompt_embeds=conditioning,
+            pooled_prompt_embeds=pooled,
+            negative_prompt_embeds=negative_conditioning,
+            negative_pooled_prompt_embeds=negative_pooled,
+            generator=generator,
+            cross_attention_kwargs=cross_attention_kwargs,
+            **model_args
+        ).images
+
+        if use_style_lora and use_char_lora:
+            st.time("unfuse lora weights")
+            self.pipe.unfuse_lora(unfuse_text_encoder=False)
+            
+        if use_style_lora or use_char_lora:
+            st.time("unload lora weights")
+            self.pipe.unload_lora_weights()
+        
+        st.time('end')
+        
+        # logger.info("Generated image: '%s'", inferences)
+        axiom_logger.info("Generated images", request_id=self.req_id, device=self.device_str, timings=st.to_str())
+        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)
+            url_name = 'https://storage.googleapis.com/' + bucket_name + '/' + image_name + '.png'
+            outputs.append(url_name)
+            axiom_logger.info("Pushed image to google cloud: "+ url_name, request_id=self.req_id, device=self.device_str)
+        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!")
+
+worker_id = os.environ.get('WORKER_ID', 'Unknown')
+if worker_id == 'Unknown':
+    raise ValueError("No worker id")
+    logger.critical("cant get worker ID")
+logger.info(f"WORKER ID: {worker_id}")
+handler = DiffusersHandler()
+handler.initialize({"gpu_id": worker_id})
+
+
+@app.route('/generate', methods=['POST'])
+def generate_image():
+    req_id = str(uuid.uuid4())
+    selected_handler = None
+    try:
+        # Extract raw requests from HTTP POST body
+        raw_requests = request.json
+        axiom_logger.info(message="Received request", request_id=req_id, **raw_requests)
+
+        gc.collect()
+        torch.cuda.empty_cache()
+        selected_handler = handler
+        selected_handler.req_id = req_id
+
+        processed_request = selected_handler.preprocess([raw_requests])
+        inferences = selected_handler.inference(processed_request)
+        outputs = selected_handler.postprocess(inferences)
+        selected_handler.req_id = None
+        return jsonify({"image_urls": outputs})
+    except Exception as e:
+        logger.error("Error during image generation: %s", str(e))
+        axiom_logger.critical("Error during image generation: " + str(e), request_id=req_id, device=selected_handler.device_str)
+        return jsonify({"error": "Failed to generate image", "details": str(e)}), 500
+
+if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=3000, threaded=False)