AWSTemplateFormatVersion: '2010-09-09'
Description: (SO8032) - Extension for Comfy on AWS - EC2 Instance
Parameters:
  StackName:
    Type: String
    Description: "The Stack Name of the ESD API"
  keyPairName:
    Description: Name of the key pair to use for the EC2 instance
    Type: AWS::EC2::KeyPair::KeyName
    ConstraintDescription: Must be the name of an existing EC2 KeyPair.
  InstanceType:
    Description: EC2 Instance Type of the Comfy to deploy
    Type: String
    AllowedValues:
      - g4dn.xlarge
      - g4dn.2xlarge
      - g4dn.4xlarge
      - g4dn.8xlarge
      - g4dn.12xlarge
      - g5.2xlarge
      - g5.4xlarge
      - g5.8xlarge
      - g5.12xlarge
    Default: g5.2xlarge
  NumberOfInferencePorts:
    Description: The number of inference ports
    Type: Number
    MinValue: 1
    MaxValue: 5
    Default: 1

Mappings:
  RegionToAmiId:
    us-east-1:
      AMI: ami-02a07d31009cc8717
    us-east-2:
      AMI: ami-02a912b010cf774bd
    ap-southeast-1:
      AMI: ami-06f0228bde7a6dac2
    ap-northeast-1:
      AMI: ami-0f9b45ca95cee6386
    ap-northeast-2:
      AMI: ami-0aba6dcbc1be4e85e
    us-west-1:
      AMI: ami-0ea5aefd1471f67b2
    us-west-2:
      AMI: ami-09d4fb52761c3d78d


Resources:
  ApiGwResource:
    Type: Custom::MyCustomResource
    Properties:
      ServiceToken: !GetAtt LambdaFunction.Arn
      StackName: !Ref StackName

  ComfyEC2Dev:
    Type: AWS::EC2::Instance
    DeletionPolicy: Delete
    Properties:
      InstanceType: !Ref InstanceType
      IamInstanceProfile: !Ref ComfyInstanceProfile
      SubnetId: !Ref ComfySubnet
      ImageId: !FindInMap [ RegionToAmiId, !Ref AWS::Region, AMI ]
      BlockDeviceMappings:
        - DeviceName: /dev/sda1
          Ebs:
            VolumeSize: 1024
            VolumeType: gp2
      KeyName: !Ref keyPairName
      SecurityGroupIds:
        - !GetAtt ComfySecurityGroup.GroupId
      Tags:
        - Key: Name
          Value: !Sub "comfy-on-aws"
      UserData:
        'Fn::Base64': !Sub |
          #!/bin/bash
          set -euxo pipefail
          
          mkdir -p ~/.aws
          echo "[default]
          region = ${AWS::Region}" > ~/.aws/config

          echo "COMFY_API_URL=${ApiGwResource.ApiGatewayUrl}" >> /etc/environment
          echo "COMFY_API_TOKEN=${ApiGwResource.ApiGatewayUrlToken}" >> /etc/environment
          echo "COMFY_BUCKET_NAME=${ApiGwResource.S3BucketName}" >> /etc/environment
          echo "DISABLE_AUTO_SYNC=false" >> /etc/environment
          echo "DISABLE_AWS_PROXY=false" >> /etc/environment
          echo "AWS_REGION=${AWS::Region}" >> /etc/environment
          echo "ACCOUNT_ID=${AWS::AccountId}" >> /etc/environment
          echo "PROCESS_NUMBER=${NumberOfInferencePorts}" >> /etc/environment
          echo "ESD_VERSION=dev" >> /etc/environment
          echo "SERVICE_TYPE=comfy" >> /etc/environment
          
          source /etc/environment

          wget https://github.com/peak/s5cmd/releases/download/v2.2.2/s5cmd_2.2.2_Linux-64bit.tar.gz
          tar xzvf s5cmd_2.2.2_Linux-64bit.tar.gz
          cp s5cmd /usr/local/bin/
          rm -rf s5cmd_2.2.2_Linux-64bit.tar.gz
          
          s5cmd cp "s3://aws-gcr-solutions-${AWS::Region}/stable-diffusion-aws-extension-github-mainline/$ESD_VERSION/src.tar" "/tmp/src.tar"
          sudo rm -rf /root/stable-diffusion-aws-extension
          sudo mkdir -p /root/stable-diffusion-aws-extension/
          tar --overwrite -xf "/tmp/src.tar" -C "/root/stable-diffusion-aws-extension/"
          rm -rf /tmp/src.tar
          
          sleep 150
          
          while sudo fuser /var/lib/apt/lists/lock >/dev/null 2>&1 ; do
            echo "Waiting for other apt process to finish..."
            sleep 1
          done
          
          while sudo fuser /var/lib/dpkg/lock-frontend >/dev/null 2>&1 ; do
            echo "Waiting for other apt process to finish..."
            sleep 1
          done
          
          sudo apt update && sudo apt install jq supervisor nginx -y
          sudo cp -f /root/stable-diffusion-aws-extension/workshop/comfy_nginx_proxy.conf /etc/nginx/sites-available/default
          sudo cp -f /root/stable-diffusion-aws-extension/workshop/comfy_custom_502.html /usr/share/nginx/html/custom_502.html
          sudo cp -f /root/stable-diffusion-aws-extension/workshop/comfy_custom_502.html /usr/share/nginx/html/index.html
          sudo ufw allow 'Nginx HTTP'
          sudo systemctl enable nginx
          sudo systemctl restart nginx

          aws ecr get-login-password --region ${AWS::Region} | docker login --username AWS --password-stdin 366590864501.dkr.ecr.${AWS::Region}.amazonaws.com
          docker pull 366590864501.dkr.ecr.${AWS::Region}.amazonaws.com/esd-inference:dev
          
          cat > comfy.service <<EOF
          [Unit]
          Description=Comfy Service
          After=network.target
          StartLimitIntervalSec=0

          [Service]
          WorkingDirectory=/root/stable-diffusion-aws-extension/
          ExecStart=bash docker_start.sh
          Type=simple
          Restart=always
          RestartSec=5
          User=root
          StartLimitAction=reboot

          [Install]
          WantedBy=default.target

          EOF
          sudo mv comfy.service /etc/systemd/system
          sudo chown root:root /etc/systemd/system/comfy.service
          
          sudo systemctl enable docker
          sudo systemctl start docker
          
          sudo systemctl enable comfy.service
          sudo systemctl start comfy.service
          
          # Step1: EC2 init
          # tail -f /var/log/cloud-init-output.log
          #
          # Step2: Comfy init
          # sudo journalctl -u comfy -f
          #
          # Step3: Comfy start
          # tail -f /root/stable-diffusion-aws-extension/container/*.log
          # 
          # sudo journalctl -u comfy --no-pager -n 200
          # docker images -q | xargs docker rmi -f



  LambdaExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: LambdaExecutionPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - logs:CreateLogGroup
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                  - cloudformation:DescribeStacks
                Resource: "*"
              - Effect: Allow
                Action:
                  - lambda:InvokeFunction
                Resource: "*"

  LambdaFunction:
    Type: AWS::Lambda::Function
    Properties:
      Handler: index.handler
      Role: !GetAtt LambdaExecutionRole.Arn
      Code:
        ZipFile: |
          # --- BEGIN PYTHON CODE ---
          import boto3
          import cfnresponse
          
          client = boto3.client('cloudformation')
          
          def handler(event, context):
            try:
              if event['RequestType'] == 'Delete':
                cfnresponse.send(event, context, cfnresponse.SUCCESS, {})
                return
          
              name = event['ResourceProperties']['StackName']
              response = client.describe_stacks(StackName=name)
              stacks = response.get('Stacks')
          
              if len(stacks) == 0:
                raise Exception('StackName Not Found')
          
              S3BucketName = ''
              ApiGatewayUrl = ''
              ApiGatewayUrlToken = ''
          
              for stack in stacks:
                stack_status = stack.get('StackStatus')
                assert stack_status in ['CREATE_COMPLETE',
                'UPDATE_COMPLETE'], 'stack is not CREATE_COMPLETE or UPDATE_COMPLETE'
                outputs = stack.get('Outputs')
                if outputs:
                  for output in outputs:
                    key = output.get('OutputKey')
                    value = output.get('OutputValue')
          
                    if key == 'S3BucketName':
                      S3BucketName = value
          
                    if key == 'ApiGatewayUrl':
                      ApiGatewayUrl = value
          
                    if key == 'ApiGatewayUrlToken':
                      ApiGatewayUrlToken = value
          
              cfnresponse.send(event, context, cfnresponse.SUCCESS, {
                'S3BucketName': S3BucketName,
                'ApiGatewayUrl': ApiGatewayUrl,
                'ApiGatewayUrlToken': ApiGatewayUrlToken,
              })
            except Exception as e:
              cfnresponse.send(event, context, cfnresponse.FAILED, {'Error': str(e)}, reason=str(e))
          
            # --- END PYTHON CODE ---
      Runtime: python3.8
      Timeout: 30

  ComfyVPC:
    Type: AWS::EC2::VPC
    DependsOn: ApiGwResource
    DeletionPolicy: Delete
    Properties:
      CidrBlock: 10.0.0.0/16
      EnableDnsSupport: true
      EnableDnsHostnames: true

  ComfySubnet:
    Type: AWS::EC2::Subnet
    DependsOn: ApiGwResource
    DeletionPolicy: Delete
    Properties:
      VpcId: !Ref ComfyVPC
      CidrBlock: 10.0.1.0/24
      MapPublicIpOnLaunch: true

  ComfyInstanceRole:
    Type: AWS::IAM::Role
    DependsOn: ApiGwResource
    DeletionPolicy: Delete
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: [ ec2.amazonaws.com ]
            Action: [ 'sts:AssumeRole' ]
      Path: "/"
      Policies:
        - PolicyName: PutObjectToSpecificBucket
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action: [ 's3:PutObject', 's3:GetObject', 's3:ListBucket', 's3:HeadObject', 's3:DeleteObject' ]
                Resource: [ '*' ]
        - PolicyName: EcrPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action: [ 'ecr:*' ]
                Resource: [ '*' ]

  ComfyInstanceProfile:
    Type: AWS::IAM::InstanceProfile
    DependsOn: ApiGwResource
    DeletionPolicy: Delete
    Properties:
      Path: "/"
      Roles:
        - !Ref ComfyInstanceRole

  ComfySecurityGroup:
    Type: AWS::EC2::SecurityGroup
    DependsOn: ApiGwResource
    DeletionPolicy: Delete
    Properties:
      VpcId: !Ref ComfyVPC
      GroupName: !Sub ${AWS::StackName}-sg
      GroupDescription: Security group for SD WebUI EC2 instance
      SecurityGroupIngress:
        - IpProtocol: tcp
          FromPort: 22
          ToPort: 22
          CidrIp: 0.0.0.0/0
        - IpProtocol: tcp
          FromPort: 10000
          ToPort: 11000
          CidrIp: 0.0.0.0/0
        - IpProtocol: tcp
          FromPort: 80
          ToPort: 80
          CidrIp: 0.0.0.0/0

  ComfyInternetGateway:
    Type: AWS::EC2::InternetGateway
    DependsOn: ApiGwResource
    DeletionPolicy: Delete

  ComfyVPCGatewayAttachment:
    Type: AWS::EC2::VPCGatewayAttachment
    DependsOn: ApiGwResource
    DeletionPolicy: Delete
    Properties:
      VpcId: !Ref ComfyVPC
      InternetGatewayId: !Ref ComfyInternetGateway

  ComfyRouteTable:
    Type: AWS::EC2::RouteTable
    DependsOn: ApiGwResource
    DeletionPolicy: Delete
    Properties:
      VpcId: !Ref ComfyVPC

  ComfyRoute:
    Type: AWS::EC2::Route
    DeletionPolicy: Delete
    Properties:
      RouteTableId: !Ref ComfyRouteTable
      DestinationCidrBlock: 0.0.0.0/0
      GatewayId: !Ref ComfyInternetGateway

  ComfySubnetRouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    DependsOn: ApiGwResource
    DeletionPolicy: Delete
    Properties:
      SubnetId: !Ref ComfySubnet
      RouteTableId: !Ref ComfyRouteTable

  ComfyEIP:
    Type: AWS::EC2::EIP
    DeletionPolicy: Delete
    Properties:
      Tags:
        - Key: Name
          Value: !Sub "comfy-on-aws"

  ComfyEIPAssociation:
    Type: AWS::EC2::EIPAssociation
    DeletionPolicy: Delete
    Properties:
      AllocationId: !GetAtt ComfyEIP.AllocationId
      InstanceId: !Ref ComfyEC2Dev

Outputs:
  TemplateSource:
    Description: Source file of the template
    Value: https://aws-gcr-solutions.s3.amazonaws.com/extension-for-stable-diffusion-on-aws/comfy_dev.yaml
  Designer:
    Description: URL for Designer
    Value: !Sub http://${ComfyEIP}
  DesignerDockerfile:
    Description: Dockerfile for EC2 and SageMaker
    Value: /root/stable-diffusion-aws-extension/container/comfy_10000.Dockerfile
  NumberOfInferencePorts:
    Description: The number of inference ports
    Value: !Sub ${NumberOfInferencePorts}
  NumberOfInferencePortsStart:
    Description: The first inference start port
    Value: !Sub http://${ComfyEIP}:10001
  PageUrlOfEC2Connect:
    Description: URL of EC2 Connect Page
    Value: !Sub https://${AWS::Region}.console.aws.amazon.com/ec2/home?region=${AWS::Region}#ConnectToInstance:instanceId=${ComfyEC2Dev}