使用生成式 AI 和 AWS 为传统媒体内容创建超分辨率

def model_fn(model_dir):
    realesr_gan_anime_model_path = os.path.join(model_dir, realesr_gan_anime_video_model_name)

    #加载RealESRGan模型
    realesr_gan_anime_video_model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=netscale, act_type='prelu')
    real_esr_gan_anime_video_upsampler = RealESRGANer(
        scale=netscale,
        model_path=realesr_gan_anime_model_path,
        dni_weight=dni_weight,
        model=realesr_gan_anime_video_model,
        tile=0,
        tile_pad=10,
        pre_pad=0,
        half=True,
        gpu_id=0)

    model = {}
    model['realesr_gan_anime'] = real_esr_gan_anime_video_upsampler
    return model

def predict_fn(input_data, model):
    # 前向传递
    input_file_path = input_data['input_file_path']
    output_file_path = input_data['output_file_path']
  
    imgname, extension = os.path.splitext(os.path.basename(input_file_path))
    img = cv2.imread(input_file_path, cv2.IMREAD_UNCHANGED)

    try:
        if ('is_anime' in input_data) and ((input_data['is_anime'].lower() == "yes") or (input_data['is_anime'].lower() == "true")):
            upsampler = model['realesr_gan_anime']
        else:
           raise RuntimeError("不支持的模型,仅支持动画上采样。")
        output, _ = upsampler.enhance(img, outscale=outscale)

    except RuntimeError as error:
        print('Error', error)
       raise error

    cv2.imwrite(output_file_path, output)
    return { "status" : 200, "output_file_path" : output_file_path, "job_id" : job_id, "batch_id" : batch_id }

def model_fn(model_dir):
    # 加载SwinIR模型
    model_path = os.path.join(model_dir, model_name)
    model = define_model(model_path, "real_sr", scale_factor)
    model = model.to(device)
    model.eval()
    return model
    
def predict_fn(input_data, model):
    # 前向传递
    input_file_path = input_data['input_file_path']
    output_file_path = input_data['output_file_path']

    img_lq = cv2.imread(input_file_path, cv2.IMREAD_COLOR).astype(np.float32) / 255.
    img_lq = np.transpose(img_lq if img_lq.shape[2] == 1 else img_lq[:, :, [2, 1, 0]], (2, 0, 1))  # HCW-BGR to CHW-RGB
    img_lq = torch.from_numpy(img_lq).float().unsqueeze(0).to(device)  # CHW-RGB to NCHW-RGB
    with torch.no_grad():
        # 将输入图像填充为窗口大小的倍数
        _, _, h_old, w_old = img_lq.size()
        h_pad = (h_old // window_size + 1) * window_size - h_old
        w_pad = (w_old // window_size + 1) * window_size - w_old
        img_lq = torch.cat([img_lq, torch.flip(img_lq, [2])], 2)[:, :, :h_old + h_pad, :]
        img_lq = torch.cat([img_lq, torch.flip(img_lq, [3])], 3)[:, :, :, :w_old + w_pad]
        output = model(img_lq)
        output = output[..., :h_old * scale_factor, :w_old * scale_factor]

        output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy()
        if output.ndim == 3:
            output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))  # CHW-RGB to HCW-BGR
        output = (output * 255.0).round().astype(np.uint8)  # float32 to uint8
        cv2.imwrite(f'{output_file_path}', output)
        torch.cuda.empty_cache()
        return { "status" : 200, "output_file_path" : output_file_path, "job_id" : job_id, "batch_id" : batch_id }

Region: <AWS区域>
Image:
  Os: alinux2
HeadNode:
  InstanceType: t3.large
  Networking:
    SubnetId: {{VPC_PUBLIC_SUBNET_ID}}
  Ssh:
    KeyName: {{SSH_KEY_PAIR_NAME}}
  CustomActions:
    OnNodeConfigured:
      Script: {{ON_NODE_STARTED_SCRIPT_HEAD_NODE}}
  Iam:
    S3Access:
      - BucketName: {{S3_BUCKET}}
        EnableWriteAccess: False
    AdditionalIamPolicies:
      - Policy: arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore
Scheduling:
  Scheduler: slurm
  SlurmSettings:
    Dns:
      UseEc2Hostnames: true
  SlurmQueues:
  - Name: q1
    ComputeResources:
    - Name: genai-api-service
      Instances:
      - InstanceType: g5.2xlarge
      MinCount: 0
      MaxCount: 1
    Networking:
      SubnetIds:
      - {{VPC_PRIVATE_SUBNET_ID}}
    CustomActions:
      OnNodeConfigured:
        Script: {{ON_NODE_STARTED_SCRIPT_GPU_COMPUTE_NODE}}
    Iam:
      S3Access:
        - BucketName: {{S3_BUCKET}}
          EnableWriteAccess: False
      AdditionalIamPolicies:
        - Policy: arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore
        - Policy: arn:aws:iam::aws:policy/EC2InstanceProfileForImageBuilderECRContainerBuilds
    Image:
      CustomAmi: {{GPU_COMPUTE_NODE_CUSTOM_AMI}}
  - Name: q2
    ComputeResources:
    - Name: video-processing
      Instances:
      - InstanceType: c6i.4xlarge
      MinCount: 1
      MaxCount: 1
    Networking:
      SubnetIds:
      - {{VPC_PRIVATE_SUBNET_ID}}
    CustomActions:
      OnNodeConfigured:
        Script: {{ON_NODE_STARTED_SCRIPT_CPU_COMPUTE_NODE}}
    Iam:
      S3Access:
        - BucketName: {{S3_BUCKET}}
          EnableWriteAccess:True
      AdditionalIamPolicies:
        - Policy: arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore
        - Policy: arn:aws:iam::aws:policy/EC2InstanceProfileForImageBuilderECRContainerBuilds
    Image:
      CustomAmi: {{CPU_COMPUTE_NODE_CUSTOM_AMI}}
SharedStorage:
  - MountDir: /fsx
    Name: fsx
    StorageType: FsxLustre
    FsxLustreSettings:
      StorageCapacity: 7200
      DeploymentType: PERSISTENT_2
      PerUnitStorageThroughput: 1000
      DeletionPolicy: Delete

pcluster create-cluster --cluster-name <集群名称> --cluster-configuration <集群配置 (cluster config) yaml>

pcluster describe-cluster --cluster-name <集群名称>

...
"cloudFormationStackStatus": "CREATE_COMPLETE", 
"clusterName": "test-cluster",
"computeFleetStatus": "STARTED",
"cloudformationStackArn": "arn:aws:cloudformation:us-east-1:123456789012:stack/test-cluster/9f2ea6a0-385e-11ee-a631-0aaaad7cc485",
"lastUpdatedTime": "2023-10-06T00:28:44.591Z",
"region": "us-east-1",
"clusterStatus": "CREATE_COMPLETE",
"scheduler": {
"type": "slurm"
}
}

#!/bin/bash
echo "downloading  bootstrap scripts from s3" >> /tmp/.node-status
aws s3 cp s3://{{S3_BUCKET_NAME}}/bootstrap/bootstrap-headnode.tar.gz /tmp/genai-video-super-resolution/
echo "completed  bootstrap scripts from s3" >> /tmp/.node-status
cd /tmp/genai-video-super-resolution
tar -xvzf bootstrap-headnode.tar.gz -C /home/ec2-user

#!/bin/bash
SECONDS=0
sudo usermod -aG docker $USER
sudo service docker restart

echo "downloading bootstrap scripts from s3" >> /tmp/.node-status
aws s3 cp s3://{{S3_BUCKET_NAME}}/bootstrap/bootstrap.tar.gz /tmp/genai-video-super-resolution-pcluster/
echo "completed bootstrap scripts from s3" >> /tmp/.node-status
cd /tmp/genai-video-super-resolution-pcluster
tar -xvzf bootstrap.tar.gz

# 将 GPU 指标添加到 cloudwatch 代理
cp cloudwatch/cloudwatch-agent-gpu-config.json /opt/aws/amazon-cloudwatch-agent/bin/config.json
/opt/aws/amazon-cloudwatch-agent/bin/amazon-cloudwatch-agent-ctl -a fetch-config -s -m ec2 -c file:/opt/aws/amazon-cloudwatch-agent/bin/config.json &>> /tmp/.node-status

# 拉取最新镜像
aws ecr get-login-password —region us-east-1 | docker login —username AWS —password-stdin 123456789012.dkr.ecr.us-east-1.amazonaws.com
echo "login successful" >> /tmp/.node-status

echo "install docker-compose" >> /tmp/.node-status
sudo curl -L https://github.com/docker/compose/releases/latest/download/docker-compose-$(uname -s)-$(uname -m) -o /usr/bin/docker-compose
sudo chmod +x /usr/bin/docker-compose
echo "docker-compose installed successfully." >> /tmp/.node-status
docker-compose -f docker/docker-compose.yaml up -d &>> /tmp/.node-status
success=0
cd test
while [ $success -ne 1 ]
do
./test.sh realesrgan &>> /tmp/.node-status
if [ $? -eq 0 ]
then
echo "done initialize realesrgan" >> /tmp/.node-status
success=1
else
echo "failed to initialize realesrgan, keep trying.." >> /tmp/.node-status
sleep 15
fi
done

success=0
while [ $success -ne 1 ]
do
./test.sh swinir2 &>> /tmp/.node-status
if [ $? -eq 0 ]
then
echo "done initialize swinir2" >> /tmp/.node-status
success=1
else
echo "failed to initialize swinir, keep trying.." >> /tmp/.node-status
sleep 15
fi
done
cd ~
echo "bootstrap took $SECONDS seconds" >> /tmp/.node-status

pcluster delete-cluster -n <集群名称> -r <AWS 区域>