feat(cosntruct): updated documentation

Author: Dinesh Sajwan

lambda/aws-contentgen-appsync-lambda/src/image_generator.py

@@ -2,7 +2,7 @@
 import requests as reqs
 from aws_lambda_powertools import Logger, Tracer, Metrics
 from requests_aws4auth import AWS4Auth
-
+from datetime import datetime
 
 logger = Logger(service="IMAGE_GENERATION")
 tracer = Tracer(service="IMAGE_GENERATION")
@@ -21,12 +21,14 @@
         )
 
 class image_generator():
-    """Generate image  ."""
+    """Generate Image based on consfigured modelId .
+        Implements text and omage moderation with Amazon Rekognition and
+        Amzon Comprehend.
+    """
 
-    def __init__(self,input_text,file_name, rekognition_client,comprehend_client,bedrock_client,bucket):
+    def __init__(self,input_text, rekognition_client,comprehend_client,bedrock_client,bucket):
                 """Initialize with bucket , key and rekognition_client."""
 
-                self.file_name = file_name
                 self.rekognition_client = rekognition_client
                 self.comprehend_client = comprehend_client
                 self.input_text =input_text
@@ -36,19 +38,21 @@ def __init__(self,input_text,file_name, rekognition_client,comprehend_client,bed
 
 
     @tracer.capture_method
-    def upload_file_to_s3(self,imgbase64encoded):
+    def upload_file_to_s3(self,imgbase64encoded,file_name):
 
         """Upload generated file to S3 bucket"""
-
-        logger.info(f"uploading file to s3 bucket: {self.bucket}, key: {self.file_name}")
+        
+        logger.info(f"uploading file to s3 bucket: {self.bucket}, key: {file_name}")
+        current_datetime = datetime.now().strftime("%Y-%m-%d_%H:%M:%S.%f")
+        upload_file_name=file_name+current_datetime+".jpg"
         try:
-            respImg= s3.Object(self.bucket, self.file_name).put(Body=base64.b64decode(imgbase64encoded))
+            respImg= s3.Object(self.bucket, upload_file_name).put(Body=base64.b64decode(imgbase64encoded))
 
         except Exception as e:
             logger.error(f"Error occured :: {e}")
             return False
         return {
-            "file_name":self.file_name,
+            "file_name":upload_file_name,
             "bucket_name":self.bucket,
         }
 
@@ -81,7 +85,7 @@ def text_moderation(self):
         return response
 
     @tracer.capture_method
-    def image_moderation(self):
+    def image_moderation(self,file_name):
 
         """Detect image moderation on the generated image to avoid any toxicity/nudity"""
 
@@ -94,7 +98,7 @@ def image_moderation(self):
         Image={
             'S3Object':{
                 'Bucket':self.bucket,
-                'Name':self.file_name}
+                'Name':file_name}
                 }
         )
         for label in rekognition_response['ModerationLabels']:
@@ -109,47 +113,33 @@ def image_moderation(self):
     def generate_image(self,input_params):
 
         """Generate image using Using bedrock with configured modelid and params"""
+        
 
         input_text=self.input_text
-        
+        print(f' input_params :: {input_params}')
         # add default negative prompts
-        if 'negative_prompts' in input_params:
+        if 'negative_prompts' in input_params and input_params['negative_prompts'] is None:
                sample_string_bytes = base64.b64decode(input_params['negative_prompts'])
                decoded_negative_prompts = sample_string_bytes.decode("utf-8")
                logger.info(f"decoded negative prompts are :: {decoded_negative_prompts}")
                negative_prompts= decoded_negative_prompts
         else:
-              negative_prompts= ["poorly rendered","poor background details"]
-    
+              negative_prompts= ["poorly rendered","poor background details","poorly drawn mountains","disfigured mountain features"]
+        
         model_id=input_params['model_config']['modelId']
 
         model_kwargs=input_params['model_config']['model_kwargs']
         params= get_inference_parameters(model_kwargs)
 
-        logger.info(f'SD params :: {params}')
 
-
-        request = json.dumps({
-                "text_prompts": (
-                        [{"text": input_text, "weight": 1.0}]
-                        + [{"text": negprompt, "weight": -1.0} for negprompt in negative_prompts]
-                ),
-                "cfg_scale":params['cfg_scale'],
-                "seed": params['seed'],
-                "steps": params['steps'],
-                "style_preset": params['style_preset'],
-                "clip_guidance_preset": params['clip_guidance_preset'],
-                "sampler": params['sampler'],
-                "width": params['width'],
-                "height": params['height']
-                })
-                        
+        body=get_model_payload(model_id,params,input_text,negative_prompts)
+        print(f' body :: {body}')
         try:
             return  self.bedrock_client.invoke_model(
                 modelId= model_id,
                 contentType= "application/json",
                 accept= "application/json",
-                body=request
+                body=body
                 )
         except Exception as e:
             logger.error(f"Error occured during generating image:: {e}")
@@ -198,7 +188,44 @@ def send_job_status(self,variables):
         )
         logger.info('res :: {}',responseJobstatus)
 
+def get_model_payload(modelid,params,input_text,negative_prompts):
+      
+     body=''
+     if modelid=='stability.stable-diffusion-xl' :
+        body = json.dumps({
+                "text_prompts": (
+                        [{"text": input_text, "weight": 1.0}]
+                        + [{"text": negprompt, "weight": -1.0} for negprompt in negative_prompts]
+                ),
+                "cfg_scale":params['cfg_scale'],
+                "seed": params['seed'],
+                "steps": params['steps'],
+                "style_preset": params['style_preset'],
+                "clip_guidance_preset": params['clip_guidance_preset'],
+                "sampler": params['sampler'],
+                "width": params['width'],
+                "height": params['height']
+                })
+        return body
+     if modelid=='amazon.titan-image-generator-v1' :
 
+        body = json.dumps({
+                       "taskType": "TEXT_IMAGE",
+                        "textToImageParams": {
+                        "text": input_text,                   
+                        #"negativeText": negative_prompts  
+                        },
+                        "imageGenerationConfig": {
+                        "numberOfImages": params['numberOfImages'],   
+                        "quality":params['quality'], 
+                        "height": params['height'],        
+                        "width": params['width'],         
+                        "cfgScale": params['cfg_scale'],      
+                        "seed": params['seed']           
+                        }
+                        })
+        return body
+      
 def get_inference_parameters(model_kwargs):
       """ Read inference parameters and set default values"""
       if 'seed' in model_kwargs:
@@ -233,6 +260,15 @@ def get_inference_parameters(model_kwargs):
               sampler= model_kwargs['sampler']
       else:
               sampler='K_DPMPP_2S_ANCESTRAL'
+      if 'numberOfImages' in model_kwargs:
+              numberOfImages= model_kwargs['numberOfImages']
+      else:
+              numberOfImages=1
+      if 'quality' in model_kwargs:
+              quality= model_kwargs['quality']
+      else:
+              quality="standard"         
+
       return {
             "cfg_scale": cfg_scale,
             "seed": seed,
@@ -242,4 +278,6 @@ def get_inference_parameters(model_kwargs):
             "sampler": sampler,
             "width": width,
             "height": height,
-      }
+            "numberOfImages": numberOfImages,
+            "quality": quality
+      }

lambda/aws-contentgen-appsync-lambda/src/lambda.py

@@ -20,7 +20,7 @@
 from aws_lambda_powertools.utilities.typing import LambdaContext
 from aws_lambda_powertools.metrics import MetricUnit
 from aws_lambda_powertools.utilities.validation import validate, SchemaValidationError
-from datetime import datetime
+
 
 
 logger = Logger(service="IMAGE_GENERATION")
@@ -30,7 +30,6 @@
 aws_region = boto3.Session().region_name
 bucket = os.environ['OUTPUT_BUCKET']
 
-
 bedrock_client = boto3.client('bedrock-runtime')
 rekognition_client=boto3.client('rekognition')
 comprehend_client=boto3.client('comprehend', region_name=aws_region)
@@ -50,17 +49,18 @@ def handler(event,  context: LambdaContext) -> dict:
     logger.info(f"event is {event}")
 
     input_params=event['detail']['imageInput']
-    #image_name="public/"+input_params['filename']
     input_text=input_params['input_text']
+    model_id=input_params['model_config']['modelId']
 
-    current_datetime = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
-    file_name="generatedimage_"+current_datetime
+    # add timestamp to file name when uploading it to s3
+    file_name="generatedimage_"
 
     response={
         "filename":file_name,
         "generatedImagePath":"",
+        "status":"Generating",
+        "image_path":bucket,
         "input_text":input_params['input_text'],
-        "generateImageJobStatus":"Failed",
         "jobid":input_params["jobid"],
         "message":''
     }
@@ -72,54 +72,74 @@ def handler(event,  context: LambdaContext) -> dict:
         response["message"]="Input text is empty."
         return response
 
-    img = image_generator(decoded_input_text,file_name, rekognition_client,comprehend_client,bedrock_client,bucket)
+    img = image_generator(decoded_input_text, rekognition_client,comprehend_client,bedrock_client,bucket)
 
     text_moderation_response=img.text_moderation()
     if(text_moderation_response['isToxic']==True):
         response["message"]="In appropriate input prompt. Please change the prompt."
+        response["status"]='Blocked'
     else:
         bedrock_response = img.generate_image(input_params)
-        parsed_reponse = parse_response(bedrock_response)
+        parsed_reponse = parse_response(bedrock_response,model_id)
         if(parsed_reponse['image_generated_status']=='Failed'):
-            response["message"]="No image generated by bedrock API, Please check the prompt"
+            response["message"]="No image generated by bedrock API, Please check the prompt."
+            response["status"]='Blocked'
         else:
-            imgbase64encoded= parsed_reponse['image_generated']
-            imageGenerated=img.upload_file_to_s3(imgbase64encoded)
-
-            image_moderation_response=img.image_moderation()
-            if(image_moderation_response['isToxic']==True):
-                response["message"]="In-appropriate image generated."
-            else:
-                response={
-                        "filename":file_name,
-                        "image_path":bucket,
-                        "input_text":decoded_input_text,
-                        "status":"Completed",
-                        "jobid":input_params["jobid"],
-                        "message":"Image generated successfully"
-                    }
-
-    print (f"response :: {response}")
-    img.send_job_status(response)
+            num_of_images=0 #if multiple image geneated iterate through all
+            for image in parsed_reponse['image_generated']:
+                print(f'num_of_images {num_of_images}')
+                if model_id=='stability.stable-diffusion-xl' :
+                    imgbase64encoded= parsed_reponse['image_generated'][num_of_images]["base64"]
+                if model_id=='amazon.titan-image-generator-v1' :
+                    imgbase64encoded= parsed_reponse['image_generated'][num_of_images]
+                imageGenerated=img.upload_file_to_s3(imgbase64encoded,file_name)
+                num_of_images=+1
+                uploaded_file_name=imageGenerated['file_name']
+                image_moderation_response=img.image_moderation(uploaded_file_name)
+                if(image_moderation_response['isToxic']==True):
+                    response["message"]="In-appropriate image generated."
+                    response["status"]='Blocked'
+                else:
+                    response={
+                            "filename":uploaded_file_name,
+                            "generatedImagePath":"",
+                            "status":"Completed",
+                            "image_path":bucket,
+                            "input_text":input_params['input_text'],
+                            "jobid":input_params["jobid"],
+                            "message":'Image generated successfully'                 
+                        }
+                print (f"response :: {response}")
+                img.send_job_status(response)
 
     return response
 
 
-def parse_response(query_response):
+def parse_response(query_response,model_id):
     """Parse response and return generated image and the prompt"""
-    print(f'query_response:: {query_response}')
-    if(not query_response):
-         parsed_reponse['image_generated_status']='Failed'
-    response_dict = json.loads(query_response["body"].read())
     parsed_reponse={
         "image_generated":'',
         "image_generated_status":'Success'
     }
-    if(response_dict['artifacts'] is None):
+    if(not query_response):
          parsed_reponse['image_generated_status']='Failed'
+         return parsed_reponse
     else:
-        parsed_reponse['image_generated']=response_dict['artifacts'][0]["base64"]
-    return parsed_reponse
+        response_dict = json.loads(query_response["body"].read())
 
+        if model_id=='stability.stable-diffusion-xl' :
 
+            if(response_dict['artifacts'] is None):
+                parsed_reponse['image_generated_status']='Failed'
+            else:
+                 parsed_reponse['image_generated']=response_dict['artifacts']
+
+        if model_id=='amazon.titan-image-generator-v1' :
+            if(response_dict['images'] is None):
+                parsed_reponse['image_generated_status']='Failed'
+            else:
+                numiofimages=response_dict['images']
+                print(f' number of images ::{len(numiofimages)}')
+                parsed_reponse['image_generated']=response_dict['images']
 
+        return parsed_reponse

lambda/aws-contentgen-appsync-lambda/src/requirements.txt

@@ -1,13 +1,8 @@
 aws-lambda-powertools
 aws-xray-sdk
 fastjsonschema
-typing-extensions
-aiohttp
 boto3>=1.28.69
 botocore>=1.31.69
 requests==2.31.0
 requests-aws4auth==1.2.3
-opensearch-py==2.3.1
-numpy
-langchain==0.0.329
-opensearch-py
+