serve loras update for logging

serve_loras.py

@@ -1,6 +1,7 @@
 from compel import Compel, ReturnedEmbeddingsType
 import logging
 from abc import ABC
+import uuid
 
 import diffusers
 import torch
@@ -24,6 +25,9 @@ import os
 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",
 )
@@ -37,6 +41,7 @@ class DiffusersHandler(ABC):
 
     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
@@ -51,6 +56,7 @@ class DiffusersHandler(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_str = device_str
         
         print("my device is " + device_str)
         self.device = torch.device(device_str)
@@ -65,6 +71,7 @@ class DiffusersHandler(ABC):
                 
         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
 
@@ -91,6 +98,7 @@ class DiffusersHandler(ABC):
         }
         
         logger.info("Processed request: '%s'", processed_request)
+        axiom_logger.info("Processed request:" + str(processed_request), request_id=self.req_id, device=self.device_str)
         return processed_request
         
 
@@ -125,6 +133,7 @@ class DiffusersHandler(ABC):
             self.pipe.unload_lora_weights()
         
         logger.info("Generated image: '%s'", inferences)
+        axiom_logger.info("Generated images", request_id=self.req_id, device=self.device_str)
         return inferences
 
     def postprocess(self, inference_outputs):
@@ -152,8 +161,9 @@ class DiffusersHandler(ABC):
             # 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')
+            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
 
 
@@ -173,6 +183,7 @@ handler_index = 0
 
 @app.route('/generate', methods=['POST'])
 def generate_image():
+    req_id = str(uuid.uuid4())
     global handler_index
     try:
         # Extract raw requests from HTTP POST body
@@ -181,14 +192,17 @@ def generate_image():
         with handler_lock:
             selected_handler = handlers[handler_index]
             handler_index = (handler_index + 1) % gpu_count  # Rotate to the next 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)
         return jsonify({"error": "Failed to generate image", "details": str(e)}), 500
 
 if __name__ == '__main__':