AWSTemplateFormatVersion: '2010-09-09'
Description: (SO8032) - Extension for Comfy on AWS - EC2 Instance
Parameters:
  Branch:
    Description: Branch of the Comfy to deploy
    Type: String
    AllowedValues:
      - main
      - dev
    Default: dev
  ApiGatewayUrl:
    Description: API URL
    Type: String
    MinLength: "1"
  ApiGatewayUrlToken:
    Description: API Key
    Type: String
    MinLength: "1"
  BucketName:
    Description: Bucket Name
    Type: String
    MinLength: "1"
  EndpointName:
    Description: Endpoint Name
    Type: String
    MinLength: "1"
  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
      - g5.2xlarge
      - g5.4xlarge
      - g5.8xlarge
    Default: g5.2xlarge
  ProcessNumber:
    Description: Number of the Comfy Processes
    Type: String
    AllowedValues:
      - 1
      - 5
      - 10
      - 15
      - 20
    Default: 10
  EsdVersion:
    Description: Version of the ESD
    Type: String
    AllowedValues:
      - latest
      - dev
    Default: latest
  WorkflowName:
    Description: Bind Workflow Name
    Type: String

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:
  ComfyVPC:
    Type: AWS::EC2::VPC
    DeletionPolicy: Delete
    Properties:
      CidrBlock: 10.0.0.0/16
      EnableDnsSupport: true
      EnableDnsHostnames: true

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

  ComfyInstanceRole:
    Type: AWS::IAM::Role
    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' ]
                Resource: [ '*' ]
        - PolicyName: EcrPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action: [ 'ecr:*' ]
                Resource: [ '*' ]

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

  ComfySecurityGroup:
    Type: AWS::EC2::SecurityGroup
    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: 8188
          ToPort: 8288
          CidrIp: 0.0.0.0/0
        - IpProtocol: tcp
          FromPort: 80
          ToPort: 80
          CidrIp: 0.0.0.0/0

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

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

  ComfyRouteTable:
    Type: AWS::EC2::RouteTable
    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
    DeletionPolicy: Delete
    Properties:
      SubnetId: !Ref ComfySubnet
      RouteTableId: !Ref ComfyRouteTable

  ComfyEC2Instance:
    Type: AWS::EC2::Instance
    DeletionPolicy: Delete
    DependsOn: ComfyInstanceProfile
    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-${Branch}" },
      ]
      UserData:
        'Fn::Base64': !Sub |
          #!/bin/bash
          set -euxo pipefail
          
          mkdir -p ~/.aws
          echo "[default]
          region = ${AWS::Region}" > ~/.aws/config

          echo "export COMFY_API_URL=${ApiGatewayUrl}" >> /etc/environment
          echo "export COMFY_API_TOKEN=${ApiGatewayUrlToken}" >> /etc/environment
          echo "export COMFY_BUCKET_NAME=${BucketName}" >> /etc/environment
          echo "export COMFY_ENDPOINT=${EndpointName}" >> /etc/environment
          echo "export DISABLE_AUTO_SYNC=false" >> /etc/environment
          echo "export DISABLE_AWS_PROXY=false" >> /etc/environment
          echo "export AWS_REGION=${AWS::Region}" >> /etc/environment
          echo "export PROCESS_NUMBER=${ProcessNumber}" >> /etc/environment
          echo "export ESD_VERSION=${EsdVersion}" >> /etc/environment
          echo "export WORKFLOW_NAME=${WorkflowName}" >> /etc/environment
          
          source /etc/environment
          
          sudo systemctl enable docker
          
          cd /root
          
          git clone https://github.com/awslabs/stable-diffusion-aws-extension.git --branch dev

          sudo apt-get update
          
          sudo apt install nginx jq -y
          sudo wget -O /etc/nginx/sites-available/default https://raw.githubusercontent.com/awslabs/stable-diffusion-aws-extension/dev/workshop/comfy_nginx_proxy.conf
          sudo wget -O /usr/share/nginx/html/custom_502.html https://raw.githubusercontent.com/awslabs/stable-diffusion-aws-extension/dev/workshop/comfy_custom_502.html
          sudo ufw allow 'Nginx HTTP'
          sudo systemctl enable nginx
          sudo systemctl restart nginx
          
          cat > comfy.service <<EOF
          [Unit]
          Description=Comfy Service
          After=network.target
          StartLimitIntervalSec=0

          [Service]
          WorkingDirectory=/root/stable-diffusion-aws-extension/
          ExecStart=bash comfy_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 comfy.service
          sudo systemctl start comfy.service

          # tail -f /var/log/cloud-init-output.log
          # sudo journalctl -u comfy -f
          # sudo journalctl -u comfy --no-pager -n 200

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

Outputs:
  ComfyNginxURL:
    Description: URL for Comfy WebUI
    Value: !Sub http://${ComfyEIP}
  ComfyRealURL:
    Description: URL for Comfy WebUI
    Value: !Sub http://${ComfyEIP}:8188
  Dockerfile:
    Description: Dockerfile for EC2 and SageMaker
    Value: /root/stable-diffusion-aws-extension/workshop/Dockerfile.comfy
  ComfyProcessStart:
    Description: Process Port Start Number
    Value: 8188