AWSTemplateFormatVersion: 2010-09-09
Description: APIGatway -> Lambda (with IAM role for apigateway, s3 and bedrock)
Parameters:
  apiGatewayName:
    Type: String
    Default: AWSBlogGetImageAPI
  apiGatewayStageName:
    Type: String
    AllowedPattern: "[a-z0-9]+"
  s3BucketName:
    Type: String
  lambdaFunctionName:
    Type: String
    Default: AWSBlogGetImage

Resources:
  ### API-GATEWAY RESOURCES ###
  PromptProcessingAPI:
    Type: 'AWS::ApiGateway::RestApi'
    Properties:
      Name: !Ref apiGatewayName
      Description: API used for invoking lambda 
      FailOnWarnings: true
      EndpointConfiguration:
        Types:
          - EDGE
  ActionResource:
    Type: 'AWS::ApiGateway::Resource'
    DependsOn:
      - PromptProcessingAPI
      - LambdaFunction
    Properties:
      RestApiId: !Ref PromptProcessingAPI
      ParentId: !GetAtt 
        - PromptProcessingAPI
        - RootResourceId
      PathPart: action
  ActionInputResource:
    Type: 'AWS::ApiGateway::Resource'
    DependsOn:
      - PromptProcessingAPI
      - LambdaFunction
      - ActionResource
    Properties:
      RestApiId: !Ref PromptProcessingAPI
      ParentId: !Ref ActionResource
      PathPart: '{actionInput}'
  PromptResource:
    Type: 'AWS::ApiGateway::Resource'
    DependsOn:
      - PromptProcessingAPI
      - LambdaFunction
      - ActionResource
      - ActionInputResource
    Properties:
      RestApiId: !Ref PromptProcessingAPI
      ParentId: !Ref ActionInputResource
      PathPart: prompt
  PromptInputResource:
    Type: 'AWS::ApiGateway::Resource'
    DependsOn:
      - PromptProcessingAPI
      - PromptResource
      - ActionResource
      - ActionInputResource
      - LambdaFunction
    Properties:
      RestApiId: !Ref PromptProcessingAPI
      ParentId: !Ref PromptResource
      PathPart: '{promptInput}'
  ProxyResource:
    Type: 'AWS::ApiGateway::Resource'
    DependsOn:
      - PromptProcessingAPI
      - PromptResource
      - ActionResource
      - ActionInputResource
      - PromptInputResource
      - LambdaFunction
    Properties:
      RestApiId: !Ref PromptProcessingAPI
      ParentId: !Ref PromptInputResource
      PathPart: '{proxy+}'
  ResponseSchema:
    Type: 'AWS::ApiGateway::Model'
    Properties:
      RestApiId: !Ref PromptProcessingAPI
      ContentType: application/json
      Description: Response Schema
      Name: ResponseSchema
      Schema:
        type: object
        required:
          - response
        properties:
          response:
            type: string
        title: Response Schema
  RequestSchema:
    Type: 'AWS::ApiGateway::Model'
    Properties:
      RestApiId: !Ref PromptProcessingAPI
      ContentType: application/json
      Description: Request Schema
      Name: RequestSchema
      Schema:
        type: object
        required:
          - request
        properties:
          request:
            type: string
        title: Request Schema
  Default4xxResponse:
    Type: AWS::ApiGateway::GatewayResponse
    Properties:
      ResponseType: DEFAULT_4XX
      RestApiId: !Ref PromptProcessingAPI
      ResponseParameters:
        gatewayresponse.header.Access-Control-Allow-Origin: "'*'"
        gatewayresponse.header.Access-Control-Allow-Headers: "'Content-Type,X-Amz-Date,Authorization,X-Api-Key,X-Amz-Security-Token'"
        gatewayresponse.header.Access-Control-Allow-Methods: "'DELETE,GET,HEAD,OPTIONS,PATCH,POST,PUT'"
        gatewayresponse.header.Access-Control-Expose-Headers: "'Date,X-Amzn-ErrorType'"
  Default5xxResponse:
    Type: AWS::ApiGateway::GatewayResponse
    Properties:
      ResponseType: DEFAULT_5XX
      RestApiId: !Ref PromptProcessingAPI
      ResponseParameters:
        gatewayresponse.header.Access-Control-Allow-Origin: "'*'"
        gatewayresponse.header.Access-Control-Allow-Headers: "'Content-Type,X-Amz-Date,Authorization,X-Api-Key,X-Amz-Security-Token'"
        gatewayresponse.header.Access-Control-Allow-Methods: "'DELETE,GET,HEAD,OPTIONS,PATCH,POST,PUT'"
        gatewayresponse.header.Access-Control-Expose-Headers: "'Date,X-Amzn-ErrorType'"
  ActionMethodCORS:
    Type: 'AWS::ApiGateway::Method'
    DependsOn:
      - PromptProcessingAPI
      - ActionResource
      - LambdaFunction
    Properties:
      RestApiId: !Ref PromptProcessingAPI
      ResourceId: !Ref ActionResource
      HttpMethod: OPTIONS
      AuthorizationType: NONE
      Integration:
        Type: MOCK
        PassthroughBehavior: WHEN_NO_MATCH
        IntegrationResponses:
          - StatusCode: 200
            ResponseParameters:
              method.response.header.Access-Control-Allow-Headers: "'Content-Type,X-Amz-Date,Authorization,X-Api-Key,X-Amz-Security-Token,X-Amz-User-Agent'"
              method.response.header.Access-Control-Allow-Methods: "'DELETE,GET,HEAD,OPTIONS,PATCH,POST,PUT'"
              method.response.header.Access-Control-Allow-Origin: "'*'"
        RequestTemplates:
          application/json: '{"statusCode": 200}'
      MethodResponses:
        - StatusCode: 200
          ResponseParameters:
            method.response.header.Access-Control-Allow-Headers: false
            method.response.header.Access-Control-Allow-Methods: false
            method.response.header.Access-Control-Allow-Origin: false
  ActionInputMethodCORS:
    Type: 'AWS::ApiGateway::Method'
    DependsOn:
      - PromptProcessingAPI
      - ActionResource
      - ActionInputResource
      - LambdaFunction
    Properties:
      RestApiId: !Ref PromptProcessingAPI
      ResourceId: !Ref ActionInputResource
      HttpMethod: OPTIONS
      AuthorizationType: NONE
      Integration:
        Type: MOCK
        PassthroughBehavior: WHEN_NO_MATCH
        IntegrationResponses:
          - StatusCode: 200
            ResponseParameters:
              method.response.header.Access-Control-Allow-Headers: "'Content-Type,X-Amz-Date,Authorization,X-Api-Key,X-Amz-Security-Token,X-Amz-User-Agent'"
              method.response.header.Access-Control-Allow-Methods: "'DELETE,GET,HEAD,OPTIONS,PATCH,POST,PUT'"
              method.response.header.Access-Control-Allow-Origin: "'*'"
        RequestTemplates:
          application/json: '{"statusCode": 200}'
      MethodResponses:
        - StatusCode: 200
          ResponseParameters:
            method.response.header.Access-Control-Allow-Headers: false
            method.response.header.Access-Control-Allow-Methods: false
            method.response.header.Access-Control-Allow-Origin: false
  ActionInputMethod:
    Type: 'AWS::ApiGateway::Method'
    DependsOn:
      - PromptProcessingAPI
      - ActionResource
      - ActionInputResource
      - LambdaFunction
    Properties:
      RestApiId: !Ref PromptProcessingAPI
      ResourceId: !Ref ActionInputResource
      HttpMethod: POST
      AuthorizationType: NONE
      RequestModels:
        application/json: !Ref RequestSchema 
      RequestParameters:
        method.request.path.actionInput: true
      Integration:
        Type: AWS_PROXY
        IntegrationHttpMethod: POST
        Uri: !Sub 
          - arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${lambdaArn}/invocations
          - lambdaArn: !GetAtt LambdaFunction.Arn
        PassthroughBehavior: WHEN_NO_MATCH
      MethodResponses:
        - StatusCode: 200
          ResponseModels: 
            application/json: !Ref ResponseSchema
  PromptMethodCORS:
    Type: 'AWS::ApiGateway::Method'
    DependsOn:
      - PromptProcessingAPI
      - PromptResource
      - LambdaFunction
    Properties:
      RestApiId: !Ref PromptProcessingAPI
      ResourceId: !Ref PromptResource
      HttpMethod: OPTIONS
      AuthorizationType: NONE
      Integration:
        Type: MOCK
        PassthroughBehavior: WHEN_NO_MATCH
        IntegrationResponses:
          - StatusCode: 200
            ResponseParameters:
              method.response.header.Access-Control-Allow-Headers: "'Content-Type,X-Amz-Date,Authorization,X-Api-Key,X-Amz-Security-Token,X-Amz-User-Agent'"
              method.response.header.Access-Control-Allow-Methods: "'DELETE,GET,HEAD,OPTIONS,PATCH,POST,PUT'"
              method.response.header.Access-Control-Allow-Origin: "'*'"
        RequestTemplates:
          application/json: '{"statusCode": 200}'
      MethodResponses:
        - StatusCode: 200
          ResponseParameters:
            method.response.header.Access-Control-Allow-Headers: false
            method.response.header.Access-Control-Allow-Methods: false
            method.response.header.Access-Control-Allow-Origin: false
  PromptInputMethodCORS:
    Type: 'AWS::ApiGateway::Method'
    DependsOn:
      - PromptProcessingAPI
      - PromptResource
      - PromptInputResource
      - LambdaFunction
    Properties:
      RestApiId: !Ref PromptProcessingAPI
      ResourceId: !Ref PromptInputResource
      HttpMethod: OPTIONS
      AuthorizationType: NONE
      Integration:
        Type: MOCK
        PassthroughBehavior: WHEN_NO_MATCH
        IntegrationResponses:
          - StatusCode: 200
            ResponseParameters:
              method.response.header.Access-Control-Allow-Headers: "'Content-Type,X-Amz-Date,Authorization,X-Api-Key,X-Amz-Security-Token,X-Amz-User-Agent'"
              method.response.header.Access-Control-Allow-Methods: "'DELETE,GET,HEAD,OPTIONS,PATCH,POST,PUT'"
              method.response.header.Access-Control-Allow-Origin: "'*'"
        RequestTemplates:
          application/json: '{"statusCode": 200}'
      MethodResponses:
        - StatusCode: 200
          ResponseParameters:
            method.response.header.Access-Control-Allow-Headers: false
            method.response.header.Access-Control-Allow-Methods: false
            method.response.header.Access-Control-Allow-Origin: false
  PromptInputMethod:
    Type: 'AWS::ApiGateway::Method'
    DependsOn:
      - PromptProcessingAPI
      - PromptResource
      - PromptInputResource
      - LambdaFunction
    Properties:
      # OperationName: processPrompt
      RestApiId: !Ref PromptProcessingAPI
      ResourceId: !Ref PromptInputResource
      HttpMethod: POST
      AuthorizationType: NONE
      RequestModels:
        application/json: !Ref RequestSchema 
      RequestParameters:
        method.request.path.actionInput: true
        method.request.path.promptInput: true
      Integration:
        Type: AWS_PROXY
        IntegrationHttpMethod: POST
        Uri: !Sub 
          - arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${lambdaArn}/invocations
          - lambdaArn: !GetAtt LambdaFunction.Arn
        PassthroughBehavior: WHEN_NO_MATCH
      MethodResponses:
        - StatusCode: 200
          ResponseModels: 
            application/json: !Ref ResponseSchema
  ProxyMethodCORS:
    Type: 'AWS::ApiGateway::Method'
    DependsOn:
      - PromptProcessingAPI
      - PromptResource
      - PromptInputResource
      - ProxyResource
      - LambdaFunction
    Properties:
      RestApiId: !Ref PromptProcessingAPI
      ResourceId: !Ref ProxyResource
      HttpMethod: OPTIONS
      AuthorizationType: NONE
      Integration:
        Type: MOCK
        PassthroughBehavior: WHEN_NO_MATCH
        IntegrationResponses:
          - StatusCode: 200
            # ResponseTemplates: {application/json: ''} ##
            ResponseParameters:
              method.response.header.Access-Control-Allow-Headers: "'Content-Type,X-Amz-Date,Authorization,X-Api-Key,X-Amz-Security-Token,X-Amz-User-Agent'"
              method.response.header.Access-Control-Allow-Methods: "'DELETE,GET,HEAD,OPTIONS,PATCH,POST,PUT'"
              method.response.header.Access-Control-Allow-Origin: "'*'"
        RequestTemplates:
          application/json: '{"statusCode": 200}'
      MethodResponses:
        - StatusCode: 200
          # ResponseModels: {application/json: Empty}
          ResponseParameters:
            method.response.header.Access-Control-Allow-Headers: false
            method.response.header.Access-Control-Allow-Methods: false
            method.response.header.Access-Control-Allow-Origin: false
  ProxyMethod:
    Type: 'AWS::ApiGateway::Method'
    DependsOn:
      - PromptProcessingAPI
      - PromptResource
      - PromptInputResource
      - ProxyResource
      - LambdaFunction
    Properties:
      RestApiId: !Ref PromptProcessingAPI
      ResourceId: !Ref ProxyResource
      HttpMethod: ANY
      AuthorizationType: NONE
      RequestModels:
        application/json: !Ref RequestSchema 
      RequestParameters:
        method.request.path.actionInput: true
        method.request.path.promptInput: true
        method.request.path.proxy: true
      Integration:
        Type: AWS_PROXY
        IntegrationHttpMethod: POST
        Uri: !Sub 
          - arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${lambdaArn}/invocations
          - lambdaArn: !GetAtt LambdaFunction.Arn
        PassthroughBehavior: WHEN_NO_MATCH
      MethodResponses:
        - StatusCode: 200
          ResponseModels: 
            application/json: !Ref ResponseSchema #{application/json: Empty}
          # ResponseParameters:
          #   method.response.header.Access-Control-Allow-Origin: false
  ApiDeployment:
    Type: AWS::ApiGateway::Deployment
    DependsOn:
      - PromptProcessingAPI
      - PromptResource
      - PromptInputResource
      - ProxyResource
      - PromptInputMethod
      - ProxyMethod
      - PromptInputMethodCORS
      - ProxyMethodCORS
      - LambdaFunction
      - Default4xxResponse
      - Default5xxResponse
    Properties:
      RestApiId: !Ref PromptProcessingAPI
      StageName: !Ref apiGatewayStageName
  LambdaFunction:
    Type: 'AWS::Lambda::Function'
    Properties:
      FunctionName: !Ref lambdaFunctionName
      Handler: index.handler
      Runtime: python3.12
      Timeout: '60'
      Role: !GetAtt LambdaExecutionRole.Arn
      ReservedConcurrentExecutions: 100
      Code:
        ZipFile: !Sub | 
          import json
          import base64
          import uuid
          import os
          import boto3
          import urllib.request
          from urllib.parse import unquote

          boto3_bedrock = boto3.client('bedrock-runtime')
          boto3_s3 = boto3.client('s3',
                                  region_name='us-east-1'
                                )
          negative_prompts = [
              "poorly rendered",
              "poor background details",
              "poorly drawn mountains",
              "disfigured mountain features",
          ]
          style_preset = "digital-art"  # (e.g. photographic, digital-art, cinematic, ...)
          clip_guidance_preset = "FAST_GREEN" # (e.g. FAST_BLUE FAST_GREEN NONE SIMPLE SLOW SLOWER SLOWEST)
          sampler = "K_DPMPP_2S_ANCESTRAL" # (e.g. DDIM, DDPM, K_DPMPP_SDE, K_DPMPP_2M, K_DPMPP_2S_ANCESTRAL, K_DPM_2, K_DPM_2_ANCESTRAL, K_EULER, K_EULER_ANCESTRAL, K_HEUN, K_LMS)
          width = 768

          def get_file_extension(file_name='png'):
              if file_name:
                  _, extension = os.path.splitext(file_name)
                  return extension.lower()
              return file_name
              
          def upload_image_to_s3(base_64_img_str,bucket_name,original_file_name=None):
              file_name = f"{original_file_name}_{uuid.uuid4().hex}.png"
              base64_bytes = base_64_img_str.encode('ascii')
              image_bytes = base64.b64decode(base64_bytes)
              boto3_s3.put_object(Bucket=bucket_name, Key=file_name, Body=image_bytes)
              return file_name                    

          def generate_image_with_bedrock( prompt, init_image_b64=None):
              ## If there is an initial image given through URL or generating based on previous output
              if init_image_b64:
                  request = json.dumps({
                      "text_prompts": (
                      [{"text": prompt, "weight": 1.0}]
                      + [{"text": negprompt, "weight": -1.0} for negprompt in negative_prompts]
                      ),
                      "cfg_scale": 10,
                      "init_image": init_image_b64,
                      "start_schedule": 0.6,
                      "steps": 50,
                      "style_preset": style_preset,
                      "clip_guidance_preset": clip_guidance_preset,
                      "sampler": sampler,
                      })
              ## If an initial image needs to be generated
              else:
                  request = json.dumps({
                      "text_prompts": (
                      [{"text": prompt, "weight": 1.0}]
                      + [{"text": negprompt, "weight": -1.0} for negprompt in negative_prompts]
                      ),
                      "cfg_scale": 5,
                      "steps": 60,
                      "style_preset": style_preset,
                      "clip_guidance_preset": clip_guidance_preset,
                      "sampler": sampler,
                      "width": width,
                      })
              modelId = "stability.stable-diffusion-xl-v1"
              response = boto3_bedrock.invoke_model(body=request, modelId=modelId)
              response_body = json.loads(response.get("body").read())
              base_64_img_str = response_body["artifacts"][0].get("base64")
              return base_64_img_str              

          def handler(event, context):
              '''
              event parameters ===>> promptInput, actionInput
              promptInput: User prompt to generate an image, e.g. Astronaut on the moon
              actionInput: Programmtic action to perform in the lambda function
                  Three types of actions:
                      1) GenerateInit: generate an initial image to begin with
                      2) s3URL: user provide s3 presigned URL and use this image to perform editing on 
                      3) UseGenerated: use generated image from instruction (1) or (2) to edit upon 
              '''
              action = event['pathParameters']['actionInput']
              prompt = event['pathParameters']['promptInput'].replace('%20', ' ')
              init_image_b64 = ''
              bucket = '${s3BucketName}'
              
              if action == 'GenerateInit':
                  base_64_img_str = generate_image_with_bedrock(prompt)
                  output_file_name = upload_image_to_s3(base_64_img_str, bucket, f'{action}_{prompt}')
                  
              else:
                  action = unquote(action)
                  action = action.split('?s3url=')
                  
                  if action[0] == 's3URL':
                      ## Image from S3 URL
                      original_img_url = action[1]
                      original_file_name = original_img_url.split('?')[0]
                      original_file_path = '/tmp/original_file_name'+get_file_extension(original_file_name)
                      
                      ## Downloading image to <original_file_path> from the link <original_img_url>
                      if original_img_url.lower().startswith('https'):
                        urllib.request.urlretrieve(original_img_url, original_file_path) 
                        with open(original_file_path, "rb") as f:
                          init_image_b64 = base64.b64encode(f.read()).decode('utf8')

                        ## Generate image
                        base_64_img_str = generate_image_with_bedrock(prompt, init_image_b64)
                        output_file_name = upload_image_to_s3(base_64_img_str, bucket, f'{action[0]}_{prompt}')
                      else:
                        output_file_name = None
                      
                  elif action[0] == 'UseGenerated':
                      ## Image from previous invoke
                      original_img_url = action[1]
                      original_file_name = original_img_url.split('?')[0]
                      original_file_path = '/tmp/original_file_name'+get_file_extension(original_file_name)
                      
                      ## Downloading image to <original_file_path> from the link <original_img_url>
                      if original_img_url.lower().startswith('https'):
                        urllib.request.urlretrieve(original_img_url, original_file_path) 
                        with open(original_file_path, "rb") as f:
                          init_image_b64 = base64.b64encode(f.read()).decode('utf8')

                        ## Generate image
                        base_64_img_str = generate_image_with_bedrock(prompt, init_image_b64)
                        output_file_name = upload_image_to_s3(base_64_img_str, bucket, f'{action[0]}_{prompt}')
                      else:
                        output_file_name = None
                  else:
                      ## Edge case handling
                      output_file_name = None
                  
              ## Generating presigned url
              if output_file_name:
                  expiration = 1800 # seconds
                  response = boto3_s3.generate_presigned_url('get_object',
                                                              Params={'Bucket': bucket,
                                                                      'Key': output_file_name},
                                                              ExpiresIn=expiration)
              else:
                  response = None
              
              output_json = {
                  'output_file_name': output_file_name,
                  'presigned_url': response
                  }
              
              return {
                  'statusCode': 200,
                  'headers': {
                      'Access-Control-Allow-Headers': 'Content-Type',
                      'Access-Control-Allow-Origin': '*',
                      'Access-Control-Allow-Methods': 'OPTIONS,POST,GET'
                  },
                  'body': json.dumps(output_json)
              }

  LambdaPermission1:
    Type: 'AWS::Lambda::Permission'
    Properties:
      Action: 'lambda:InvokeFunction'
      FunctionName: !GetAtt LambdaFunction.Arn
      Principal: apigateway.amazonaws.com
      SourceArn: !Sub arn:aws:execute-api:${AWS::Region}:${AWS::AccountId}:${PromptProcessingAPI}/*/POST/action/*/prompt/*

  LambdaPermission2:
    Type: 'AWS::Lambda::Permission'
    Properties:
      Action: 'lambda:InvokeFunction'
      FunctionName: !GetAtt LambdaFunction.Arn
      Principal: apigateway.amazonaws.com
      SourceArn: !Sub arn:aws:execute-api:${AWS::Region}:${AWS::AccountId}:${PromptProcessingAPI}/*/POST/action/*

  LambdaPermission3:
    Type: 'AWS::Lambda::Permission'
    Properties:
      Action: 'lambda:InvokeFunction'
      FunctionName: !GetAtt LambdaFunction.Arn
      Principal: apigateway.amazonaws.com
      SourceArn: !Sub arn:aws:execute-api:${AWS::Region}:${AWS::AccountId}:${PromptProcessingAPI}/*/*/action/*/prompt/*/*
  
  LambdaExecutionRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      Policies:
        - PolicyName: lambda
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action: 
                  - 'logs:CreateLogGroup'
                  - 'logs:CreateLogStream'
                  - 'logs:PutLogEvents'
                  - 'logs:DescribeLogStreams'
                Resource:
                  - !Sub arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/lambda/${lambdaFunctionName}:*
        - PolicyName: bedrock-invoke-access
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - 'bedrock:InvokeModel'
                Resource: !Sub 'arn:aws:bedrock:${AWS::Region}::foundation-model/stability.stable-diffusion-xl-v1'
        - PolicyName: s3-read-write-restricted
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - 's3:ListBucket'
                Resource: !Sub 'arn:aws:s3:::${s3BucketName}'
              - Effect: Allow
                Action:
                  - 's3:PutObject'
                  - 's3:GetObject'
                  - 's3:DeleteObject'
                  - 's3:ListBucket'
                Resource: !Sub 'arn:aws:s3:::${s3BucketName}/*'

  lambdaLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Sub /aws/lambda/${lambdaFunctionName}
      RetentionInDays: 30

Outputs:
  apiGatewayInvokeURL:
    Value: !Sub https://${PromptProcessingAPI}.execute-api.${AWS::Region}.amazonaws.com/${apiGatewayStageName}