Feature added an Option for Either Generating an S3 url or b64 json

product_diffusion_api/routers/sdxl_text_to_image.py

@@ -1,5 +1,6 @@
 import sys
 from torchao.quantization import apply_dynamic_quant
+
 sys.path.append("../scripts")  # Path of the scripts directory
 import config
 from fastapi import APIRouter, HTTPException
@@ -11,6 +12,9 @@ import uuid
 from diffusers import DiffusionPipeline
 import torch
 from functools import lru_cache
+from s3_manager import S3ManagerService
+from PIL import Image
+import io
 
 torch._inductor.config.conv_1x1_as_mm = True
 torch._inductor.config.coordinate_descent_tuning = True
@@ -52,18 +56,30 @@ def dynamic_quant_filter_fn(mod, *args):
 
 
 
-# Utility function to convert PIL image to base64 encoded JSON
+
+
 def pil_to_b64_json(image):
-    # Generate a UUID for the image
     image_id = str(uuid.uuid4())
     buffered = BytesIO()
     image.save(buffered, format="PNG")
     b64_image = base64.b64encode(buffered.getvalue()).decode("utf-8")
     return {"image_id": image_id, "b64_image": b64_image}
 
-def upload_pil_to_s3(image):
-    image
-    
+
+def pil_to_s3_json(image: Image.Image, file_name: str) -> str:
+    image_id = str(uuid.uuid4())
+    s3_uploader = S3ManagerService()
+    image_bytes = io.BytesIO()
+    image.save(image_bytes, format="PNG")
+    image_bytes.seek(0)
+
+    unique_file_name = s3_uploader.generate_unique_file_name(file_name)
+    s3_uploader.upload_file(image_bytes, unique_file_name)
+    signed_url = s3_uploader.generate_signed_url(
+        unique_file_name, exp=43200
+    )  # 12 hours
+    return {"image_id": image_id, "url": signed_url}
+
 
 @lru_cache(maxsize=1)
 def load_pipeline(model_name, adapter_name):
@@ -77,9 +93,6 @@ def load_pipeline(model_name, adapter_name):
     pipe.fuse_qkv_projections()
     apply_dynamic_quant(pipe.unet, dynamic_quant_filter_fn)
     apply_dynamic_quant(pipe.vae, dynamic_quant_filter_fn)
-    
-  
-
     return pipe
 
 
@@ -125,12 +138,17 @@ class SDXLLoraInference:
         self.num_inference_steps = num_inference_steps
         self.guidance_scale = guidance_scale
 
-    def run_inference(self):
+    def run_inference(self, mode: str = "b64_json") -> str:
         """
         Runs the inference process and returns the generated image.
 
+        Parameters:
+            mode (str): The mode for returning the generated image.
+                        Possible values: "b64_json", "s3_json".
+                        Defaults to "b64_json".
+
         Returns:
-            str: The generated image in base64-encoded JSON format.
+            str: The generated image in the specified format.
         """
         image = self.pipe(
             prompt=self.prompt,
@@ -139,7 +157,14 @@ class SDXLLoraInference:
             negative_prompt=self.negative_prompt,
             num_images_per_prompt=self.num_images,
         ).images[0]
-        return pil_to_b64_json(image)
+        
+        if mode == "s3_json":
+            s3_url = pil_to_s3_json(image)
+            return pil_to_s3_json(image, s3_url)
+        elif mode == "b64_json":
+            return pil_to_b64_json(image)
+        else:
+            raise ValueError("Invalid mode. Supported modes are 'b64_json' and 's3_json'.")
 
 
 # Input format for single request
@@ -170,7 +195,7 @@ async def sdxl_v0_lora_inference(data: InputFormat):
     return output_json
 
 
-# Endpoint for batch requests
+
 @router.post("/sdxl_v0_lora_inference/batch")
 async def sdxl_v0_lora_inference_batch(data: BatchInputFormat):
     """

scripts/s3_manager.py

@@ -1,33 +1,33 @@
 import base64
 import io
+import os
 import boto3
 from botocore.config import Config
 import random
 import string
 from dotenv import load_dotenv
 
-env = load_dotenv('../config.env')
+# Load environment variables from the .env file
+load_dotenv('../config.env')
 
-
-class ImageService:
-    def _init_(self):
+class S3ManagerService:
+    def __init__(self):
         self.s3 = boto3.client(
             "s3",
             config=Config(signature_version="s3v4"),
-            aws_access_key_id=env.AWS_ACCESS_KEY_ID,
-            aws_secret_access_key=env.AWS_SECRET_ACCESS_KEY,
-            region_name=env.AWS_REGION,
+            aws_access_key_id=os.getenv('AWS_ACCESS_KEY_ID'),
+            aws_secret_access_key=os.getenv('AWS_SECRET_ACCESS_KEY'),
+            region_name=os.getenv('AWS_REGION'),
         )
-
-    def generate_signed_url(self, file_name: str, exp: int = 1800) -> str:
+    
+    def generate_signed_url(self, file_name: str, exp: int = 43200) -> str:  # 43200 seconds = 12 hours
         return self.s3.generate_presigned_url(
             "get_object",
-            Params={"Bucket": env.AWS_BUCKET_NAME, "Key": file_name},
+            Params={"Bucket": os.getenv('AWS_BUCKET_NAME'), "Key": file_name},
             ExpiresIn=exp,
         )
 
-    def generate_unique_file_name(self, file) -> str:
-        file_name = file.filename
+    def generate_unique_file_name(self, file_name: str) -> str:
         random_string = "".join(
             random.choices(string.ascii_uppercase + string.digits, k=10)
         )
@@ -36,7 +36,7 @@ class ImageService:
         return f"{file_real_name}-{random_string}.{file_extension}"
 
     def upload_file(self, file, file_name) -> str:
-        self.s3.upload_fileobj(file, env.AWS_BUCKET_NAME, file_name)
+        self.s3.upload_fileobj(file, os.getenv('AWS_BUCKET_NAME'), file_name)
         return file_name
 
     def upload_base64_file(self, base64_file: str, file_name: str) -> str:
@@ -45,5 +45,9 @@ class ImageService:
     def get_object(self, file_name: str, bucket: str):
         try:
             return self.s3.get_object(Bucket=bucket, Key=file_name)
-        except:  # noqa: E722
+        except self.s3.exceptions.NoSuchKey:
+            print(f"The file {file_name} does not exist in the bucket {bucket}.")
+            return None
+        except Exception as e:
+            print(f"An error occurred: {e}")
             return None