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

Mappings:
  RegionToAmiId:
    us-east-1:
      AMI: ami-0e4cc88070dfcf17b
    us-east-2:
      AMI: ami-0a48cc2a4a0988c70
    ap-southeast-1:
      AMI: ami-09a7fb29eb7375468
    ap-northeast-1:
      AMI: ami-0a1f6795aeb907a68

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' ]
                Resource: [ !Sub 'arn:aws:s3:::${BucketName}/*', !Sub 'arn:aws:s3:::${BucketName}' ]

  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: 8188
          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
    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: gp3
      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=true" >> /etc/environment
          echo "export DISABLE_AWS_PROXY=false" >> /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/

          sudo apt-get update
          
          sudo apt install nginx -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
          
          sudo add-apt-repository ppa:deadsnakes/ppa -y
          sudo apt update -y
          sudo apt --fix-broken install wget python3.10 python3.10-venv python3-dev python3.10-dev python3-pip git libgl1 libtcmalloc-minimal4 -y
          sudo update-alternatives --install /usr/bin/python3 python /usr/bin/python3.10 1
          curl -sS https://bootstrap.pypa.io/get-pip.py | sudo python3.10
          sudo python3 -m venv venv
          source venv/bin/activate
          
          sudo curl -sSL "https://raw.githubusercontent.com/awslabs/stable-diffusion-aws-extension/dev/build_scripts/install_comfy.sh" | sudo bash;
          sudo rm /ComfyUI/custom_nodes/comfy_sagemaker_proxy.py
          
          sudo mkdir -p /ComfyUI/models/vae/
          sudo wget -O /ComfyUI/models/vae/vae-ft-mse-840000-ema-pruned.safetensors https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.safetensors
          
          sudo mkdir -p /ComfyUI/models/checkpoints/
          sudo wget -O /ComfyUI/models/checkpoints/majicmixRealistic_v7.safetensors https://huggingface.co/GreenGrape/231209/resolve/045ebfc504c47ba8ccc424f1869c65a223d1f5cc/majicmixRealistic_v7.safetensors
          
          sudo mkdir -p /ComfyUI/models/animatediff_models/
          sudo wget -O /ComfyUI/models/animatediff_models/mm_sd_v15_v2.ckpt https://huggingface.co/guoyww/animatediff/resolve/main/mm_sd_v15_v2.ckpt
          sudo chmod -R 777 /ComfyUI
          
          cat > comfy.service <<EOF
          [Unit]
          Description=Comfy Service
          After=network.target
          StartLimitIntervalSec=0

          [Service]
          WorkingDirectory=/ComfyUI
          ExecStart=sudo /venv/bin/python3 main.py --listen
          Type=simple
          Restart=always
          RestartSec=3
          User=ubuntu
          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

#          cat > comfy_upload.service <<EOF
#          [Unit]
#          Description=Comfy Upload Service
#          After=network.target
#          StartLimitIntervalSec=0
#
#          [Service]
#          WorkingDirectory=/ComfyUI
#          ExecStart=/stable-diffusion-aws-extension/build_scripts/comfy/inotifywait.sh
#          Type=simple
#          Restart=always
#          RestartSec=3
#          User=root
#          StartLimitAction=reboot
#
#          [Install]
#          WantedBy=default.target
#
#          EOF
#          sudo mv comfy_upload.service /etc/systemd/system
#          sudo chown root:root /etc/systemd/system/comfy_upload.service
#          sudo systemctl enable comfy_upload.service
#          sudo systemctl start comfy_upload.service

          # tail -f /var/log/cloud-init-output.log
          # sudo journalctl -u comfy -f
          # sudo journalctl -u comfy_upload -f
          # sudo journalctl -u comfy --no-pager -n 200
          # wget -O /ComfyUI/models/checkpoints/v1-5-pruned-emaonly.ckpt https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.ckpt?download=true

  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
