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-2:
      AMI: ami-063b634d62a631728
    ap-south-1:
      AMI: ami-007020fd9c84e18c7
    eu-south-1:
      AMI: ami-0fdd5958b81f36e68
    eu-south-2:
      AMI: ami-0ee19dcf57f09e938
    me-central-1:
      AMI: ami-04be4fee4051bf815
    ca-central-1:
      AMI: ami-036149658febd9fd6
    eu-central-1:
      AMI: ami-0b584e86060d6e0f6
    eu-central-2:
      AMI: ami-059f62f210ee212db
    us-west-1:
      AMI: ami-036cafe742923b3d9
    us-west-2:
      AMI: ami-073ff6027d02b1312
    af-south-1:
      AMI: ami-0a7945c9f7659f25f
    eu-north-1:
      AMI: ami-0a1df6b0891ee52a7
    eu-west-3:
      AMI: ami-09e513e9eacab10c1
    eu-west-2:
      AMI: ami-09627c82937ccdd6d
    eu-west-1:
      AMI: ami-0d0fa503c811361ab
    ap-northeast-3:
      AMI: ami-08466eb240ce14dfe
    ap-northeast-2:
      AMI: ami-0fb936d7ba52ec375
    me-south-1:
      AMI: ami-0e0aab5c81dcbb143
    ap-northeast-1:
      AMI: ami-09a81b370b76de6a2
    sa-east-1:
      AMI: ami-07871d79cd57217f5
    ap-east-1:
      AMI: ami-04e13bca8e679e5ce
    ap-southeast-1:
      AMI: ami-02ee763250491e04a
    ap-southeast-2:
      AMI: ami-076fe60835f136dc9
    ap-southeast-3:
      AMI: ami-0913922d1289852b6
    ap-southeast-4:
      AMI: ami-03959cf65827f9f43
    us-east-1:
      AMI: ami-051dcca84f1edfff1
    us-east-2:
      AMI: ami-065d315e052507855

Resources:
  WebUiEC2Instance:
    Type: AWS::EC2::Instance
    DeletionPolicy: Delete
    DependsOn: WebUiInstanceProfile
    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" },
      ]
      UserData:
        'Fn::Base64': !Sub |
          #!/bin/bash
          set -euxo pipefail
          
          mkdir -p ~/.aws
          echo "[default]
          region = ${AWS::Region}" > ~/.aws/config
          
          echo "AWS_REGION=${AWS::Region}" >> /etc/environment
          echo "ACCOUNT_ID=${AWS::AccountId}" >> /etc/environment
          echo "AWS_DEFAULT_REGION=${AWS::Region}" >> /etc/environment
          echo "ESD_EC2=true" >> /etc/environment
          echo "ESD_VERSION=dev" >> /etc/environment
          source /etc/environment
          
          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-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/sd_nginx_proxy.conf
          sudo wget -O /usr/share/nginx/html/custom_502.html https://raw.githubusercontent.com/awslabs/stable-diffusion-aws-extension/dev/workshop/sd_custom_502.html
          sudo wget -O /usr/share/nginx/html/index.html https://raw.githubusercontent.com/awslabs/stable-diffusion-aws-extension/dev/workshop/sd_custom_502.html
          sudo ufw allow 'Nginx HTTP'
          sudo systemctl enable nginx
          sudo systemctl restart nginx
          
          sudo apt install software-properties-common -y
          sudo apt install wget git python3.10 python3.10-venv python3-dev 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/dev/build_scripts/install_sd.sh" | bash;
          
          cd stable-diffusion-webui/extensions
          git clone https://github.com/zixaphir/Stable-Diffusion-Webui-Civitai-Helper.git
          
          cd stable-diffusion-aws-extension

          source /home/ubuntu/stable-diffusion-webui/venv/bin/activate
          pip install tensorboard==2.15.2 tensorflow==2.15.0.post1
          pip install ml_dtypes==0.4.0
          pip install boto3

          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/
          
          bucket="aws-gcr-solutions-$AWS_REGION"
          prefix="stable-diffusion-aws-extension-github-mainline/models"
          echo "cp s3://$bucket/$prefix/sd_xl_base_1.0.safetensors /home/ubuntu/stable-diffusion-webui/models/Stable-diffusion/" > /tmp/models.txt
          echo "cp s3://$bucket/$prefix/v1-5-pruned-emaonly.safetensors /home/ubuntu/stable-diffusion-webui/models/Stable-diffusion/" >> /tmp/models.txt
          echo "cp s3://$bucket/$prefix/LahCuteCartoonSDXL_alpha.safetensors /home/ubuntu/stable-diffusion-webui/models/Stable-diffusion/" >> /tmp/models.txt
          echo "cp s3://$bucket/$prefix/majicmixRealistic_v7.safetensors /home/ubuntu/stable-diffusion-webui/models/Stable-diffusion/" >> /tmp/models.txt
          echo "cp s3://$bucket/$prefix/control_v11p_sd15_canny.pth /home/ubuntu/stable-diffusion-webui/models/ControlNet/" >> /tmp/models.txt
          echo "cp s3://$bucket/$prefix/control_v11p_sd15_openpose.pth /home/ubuntu/stable-diffusion-webui/models/ControlNet/" >> /tmp/models.txt
          echo "cp s3://$bucket/$prefix/lcm_lora_xl.safetensors /home/ubuntu/stable-diffusion-webui/models/Lora/" >> /tmp/models.txt
          echo "cp s3://$bucket/$prefix/lcm_lora_1_5.safetensors /home/ubuntu/stable-diffusion-webui/models/Lora/" >> /tmp/models.txt
          echo "cp s3://$bucket/$prefix/nendoroid_xl_v7.safetensors /home/ubuntu/stable-diffusion-webui/models/Lora/" >> /tmp/models.txt
          s5cmd run /tmp/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
          
          # tail -f /var/log/cloud-init-output.log
          # sudo journalctl -u sd-webui -f


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

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

  WebUiInstanceRole:
    Type: AWS::IAM::Role
    DeletionPolicy: Delete
    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
        - arn:aws:iam::aws:policy/AmazonEC2ReadOnlyAccess

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

  WebUiSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    DeletionPolicy: Delete
    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
        - IpProtocol: tcp
          FromPort: 6006
          ToPort: 6006
          CidrIp: 0.0.0.0/0

  WebUiInternetGateway:
    Type: AWS::EC2::InternetGateway
    DeletionPolicy: Delete

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

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

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

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

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

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

Outputs:
  TemplateSource:
    Description: Source file of the template
    Value: https://aws-gcr-solutions.s3.amazonaws.com/extension-for-stable-diffusion-on-aws/sd_dev.yaml
  WebUINginxURL:
    Description: URL for Stable Diffusion WebUI
    Value: !Sub http://${WebUiEIP}
  WebUIRealURL:
    Description: URL for Stable Diffusion WebUI
    # add port 7860 to the end of the URL
    Value: !Sub http://${WebUiEIP}:7860
  PageUrlOfEC2Connect:
    Description: URL of EC2 Connect Page
    Value: !Sub https://${AWS::Region}.console.aws.amazon.com/ec2/home?region=${AWS::Region}#ConnectToInstance:instanceId=${WebUiEC2Instance}