AWSTemplateFormatVersion: '2010-09-09'
Description: (SO8032) - Extension for Stable Diffusion on AWS - EC2 Instance
Parameters:
  Branch:
    Description: Branch of the ESD to deploy
    Type: String
    AllowedValues:
      - main
      - dev
    Default: main
  InstanceType:
    Description: EC2 Instance Type of the ESD to deploy
    Type: String
    AllowedValues:
      - c5.2xlarge
      - c5.4xlarge
      - g5.2xlarge
      - g5.4xlarge
    Default: c5.4xlarge

Mappings:
  RegionToAmiId:
    ap-south-1:
      AMI: ami-03f4878755434977f
    eu-north-1:
      AMI: ami-0014ce3e52359afbd
    eu-west-3:
      AMI: ami-01d21b7be69801c2f
    eu-west-2:
      AMI: ami-0e5f882be1900e43b
    eu-west-1:
      AMI: ami-0905a3c97561e0b69
    ap-northeast-3:
      AMI: ami-05ff0b3a7128cd6f8
    ap-northeast-2:
      AMI: ami-0f3a440bbcff3d043
    ap-northeast-1:
      AMI: ami-07c589821f2b353aa
    ca-central-1:
      AMI: ami-0a2e7efb4257c0907
    sa-east-1:
      AMI: ami-0fb4cf3a99aa89f72
    ap-east-1:
      AMI: ami-0d96ec8a788679eb2
    ap-southeast-1:
      AMI: ami-0fa377108253bf620
    ap-southeast-2:
      AMI: ami-04f5097681773b989
    eu-central-1:
      AMI: ami-0faab6bdbac9486fb
    us-east-1:
      AMI: ami-0c7217cdde317cfec
    us-east-2:
      AMI: ami-05fb0b8c1424f266b
    us-west-1:
      AMI: ami-0ce2cb35386fc22e9
    us-west-2:
      AMI: ami-008fe2fc65df48dac

Resources:
  WebUiVPC:
    Type: AWS::EC2::VPC
    Properties:
      CidrBlock: 10.0.0.0/16
      EnableDnsSupport: true
      EnableDnsHostnames: true

  WebUiSubnet:
    Type: AWS::EC2::Subnet
    Properties:
      VpcId: !Ref WebUiVPC
      CidrBlock: 10.0.1.0/24
      MapPublicIpOnLaunch: true

  WebUiInstanceRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: [ ec2.amazonaws.com ]
            Action: [ 'sts:AssumeRole' ]
      Path: "/"
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/AmazonS3ReadOnlyAccess

  WebUiInstanceProfile:
    Type: AWS::IAM::InstanceProfile
    Properties:
      Path: "/"
      Roles:
        - Ref: WebUiInstanceRole

  WebUiSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      VpcId: !Ref WebUiVPC
      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: 80
          ToPort: 80
          CidrIp: 0.0.0.0/0
        - IpProtocol: tcp
          FromPort: 7860
          ToPort: 7860
          CidrIp: 0.0.0.0/0

  WebUiInternetGateway:
    Type: AWS::EC2::InternetGateway

  WebUiVPCGatewayAttachment:
    Type: AWS::EC2::VPCGatewayAttachment
    Properties:
      VpcId: !Ref WebUiVPC
      InternetGatewayId: !Ref WebUiInternetGateway

  WebUiRouteTable:
    Type: AWS::EC2::RouteTable
    Properties:
      VpcId: !Ref WebUiVPC

  WebUiRoute:
    Type: AWS::EC2::Route
    Properties:
      RouteTableId: !Ref WebUiRouteTable
      DestinationCidrBlock: 0.0.0.0/0
      GatewayId: !Ref WebUiInternetGateway

  WebUiSubnetRouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      SubnetId: !Ref WebUiSubnet
      RouteTableId: !Ref WebUiRouteTable

  WebUiEC2Instance:
    Type: AWS::EC2::Instance
    Properties:
      InstanceType: !Ref InstanceType
      IamInstanceProfile: !Ref WebUiInstanceProfile
      SubnetId: !Ref WebUiSubnet
      ImageId: !FindInMap [ RegionToAmiId, !Ref AWS::Region, AMI ]
      BlockDeviceMappings:
        - DeviceName: /dev/sda1
          Ebs:
            VolumeSize: 1024
            VolumeType: gp3
      SecurityGroupIds:
        - !GetAtt WebUiSecurityGroup.GroupId
      Tags: [
        { "Key": "Name", "Value": !Sub "sd-on-aws-${Branch}" },
      ]
      UserData:
        'Fn::Base64': !Sub |
          #!/bin/bash
          set -euxo pipefail
          
          sudo apt-get update
          
          sudo apt install nginx -y
          sudo ufw allow 'Nginx HTTP'
          sudo wget -O /etc/nginx/sites-available/default https://raw.githubusercontent.com/awslabs/stable-diffusion-aws-extension/dev/workshop/nginx_proxy.conf
          sudo wget -O /usr/share/nginx/html/custom_502.html https://raw.githubusercontent.com/awslabs/stable-diffusion-aws-extension/dev/workshop/custom_502.html
          sudo systemctl enable nginx
          sudo systemctl restart nginx
          
          sudo apt install software-properties-common -y
          sudo add-apt-repository ppa:deadsnakes/ppa -y
          sudo apt install wget git python3.10 python3.10-venv build-essential net-tools libgl1 libtcmalloc-minimal4 -y
          sudo update-alternatives --install /usr/bin/python3 python /usr/bin/python3.10 1
          
          cd /home/ubuntu
          
          curl -sSL https://raw.githubusercontent.com/awslabs/stable-diffusion-aws-extension/${Branch}/install.sh | bash
          
          cd stable-diffusion-webui/extensions
          git clone https://github.com/zixaphir/Stable-Diffusion-Webui-Civitai-Helper.git
          
          cd stable-diffusion-aws-extension
          git checkout ${Branch}

          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
          export AWS_REGION=us-east-1
          wget https://raw.githubusercontent.com/awslabs/stable-diffusion-aws-extension/dev/workshop/models.txt
          sudo mv s5cmd /usr/bin/
          s5cmd run models.txt
          
          sudo chown -R ubuntu:ubuntu /home/ubuntu/stable-diffusion-webui
          
          cat > sd-webui.service <<EOF
          [Unit]
          Description=Stable Diffusion UI server
          After=network.target
          StartLimitIntervalSec=0

          [Service]
          WorkingDirectory=/home/ubuntu/stable-diffusion-webui
          ExecStart=/home/ubuntu/stable-diffusion-webui/webui.sh --enable-insecure-extension-access --skip-torch-cuda-test --no-half --listen --no-download-sd-model
          Type=simple
          Restart=always
          RestartSec=3
          User=ubuntu
          StartLimitAction=reboot

          [Install]
          WantedBy=default.target

          EOF
          sudo mv sd-webui.service /etc/systemd/system
          sudo chown root:root /etc/systemd/system/sd-webui.service
          sudo systemctl enable sd-webui.service
          sudo systemctl start sd-webui.service

  WebUiEIP:
    Type: AWS::EC2::EIP
  WebUiEIPAssociation:
    Type: AWS::EC2::EIPAssociation
    Properties:
      AllocationId: !GetAtt WebUiEIP.AllocationId
      InstanceId: !Ref WebUiEC2Instance

Outputs:
  WebUIURL:
    Description: URL for Stable Diffusion WebUI
    # add port 7860 to the end of the URL
    Value: !Sub http://${WebUiEIP}