merge proxy

.gitignore

@@ -48,5 +48,5 @@ test/aigc_webui_inference_images/.env
 test/**/.env
 *.iml
 .DS_Store
-/workshop/ComfyUI/
+/ComfyUI/
 /.env

Dockerfile.comfy

@@ -0,0 +1,10 @@
+ARG AWS_REGION
+FROM 366590864501.dkr.ecr.$AWS_REGION.amazonaws.com/esd-inference:dev
+
+# TODO BYOC
+#RUN apt-get update -y && \
+#    apt-get install ffmpeg -y && \
+#    rm -rf /var/lib/apt/lists/* \
+
+COPY build_scripts/inference/start.sh /
+RUN chmod +x /start.sh

build_scripts/inference/start.sh

@@ -280,7 +280,6 @@ comfy_launch(){
   chmod -R +x venv/bin
 
   rm -rf /home/ubuntu/ComfyUI/custom_nodes/ComfyUI-AWS-Extension
-  rm /home/ubuntu/ComfyUI/custom_nodes/comfy_local_proxy.py
   source venv/bin/activate
   python /metrics.py &
 
@@ -331,29 +330,12 @@ comfy_launch_from_public_s3(){
 
 # -------------------- startup --------------------
 
-# if pipeline finished, it will be executed
-#if [[ $IMAGE_URL == *"dev"* ]]; then
-#  download_conda
-#  if [ "$SERVICE_TYPE" == "sd" ]; then
-#      sd_install_build
-#      /serve trim_sd.sh
-#      sd_cache_endpoint
-#      sd_launch
-#      exit 1
-#  else
-#      comfy_install_build
-#      /serve trim_comfy
-#      comfy_cache_endpoint
-#      comfy_launch
-#      exit 1
-#  fi
-#fi
-
 ec2_start_process(){
   set -euxo pipefail
   echo "---------------------------------------------------------------------------------"
   export LD_LIBRARY_PATH=$LD_PRELOAD
   download_conda
+
   init_port=8187
   for i in $(seq 1 "$PROCESS_NUMBER"); do
       init_port=$((init_port + 1))
@@ -381,7 +363,38 @@ ec2_start_process(){
   done
 }
 
-if [ -n "$COMFY_EC2" ]; then
+if [ -n "$WORKFLOW_NAME" ]; then
+  start_at=$(date +%s)
+  s5cmd --log=error sync "s3://$S3_BUCKET_NAME/comfy/workflows/$WORKFLOW_NAME/*" "/home/ubuntu/ComfyUI/"
+  end_at=$(date +%s)
+  export DOWNLOAD_FILE_SECONDS=$((end_at-start_at))
+  echo "download file: $DOWNLOAD_FILE_SECONDS seconds"
+
+  cd "/home/ubuntu/ComfyUI" || exit 1
+
+  rm -rf web/extensions/ComfyLiterals
+
+  chmod -R 777 "/home/ubuntu/ComfyUI"
+  chmod -R +x venv
+  source venv/bin/activate
+
+  # on EC2
+  if [ -n "$ON_EC2" ]; then
+    ec2_start_process
+    exit 1
+  fi
+
+  if [ -n "$ON_SAGEMAKER" ]; then
+    python3 serve.py
+    exit 1
+  fi
+
+  # on SageMaker
+  python3 serve.py
+  exit 1
+fi
+
+if [ -n "$ON_EC2" ]; then
   set -euxo pipefail
   cd /home/ubuntu || exit 1
 
@@ -409,8 +422,6 @@ if [ -n "$COMFY_EC2" ]; then
   export DECOMPRESS_SECONDS=$((end_at-start_at))
   echo "decompress file: $DECOMPRESS_SECONDS seconds"
 
-  ls -la
-
   rm ./ComfyUI/custom_nodes/comfy_sagemaker_proxy.py
 
   cd /home/ubuntu/ComfyUI || exit 1
@@ -439,28 +450,6 @@ if [ -n "$COMFY_EC2" ]; then
   exit 1
 fi
 
-if [ -n "$APP_SOURCE" ]; then
-  if [ -n "$APP_CWD" ]; then
-    start_at=$(date +%s)
-    s5cmd --log=error sync "s3://$S3_BUCKET_NAME/${APP_SOURCE}*" "$APP_CWD"
-    end_at=$(date +%s)
-    export DOWNLOAD_FILE_SECONDS=$((end_at-start_at))
-    echo "download file: $DOWNLOAD_FILE_SECONDS seconds"
-
-    cd "$APP_CWD" || exit 1
-
-    rm -rf web/extensions/ComfyLiterals
-
-    chmod -R 777 "$APP_CWD"
-    chmod -R +x venv
-
-    source venv/bin/activate
-
-    python3 serve.py
-
-    exit 1
-  fi
-fi
 
 if [ -f "/initiated_lock" ]; then
     echo "already initiated, start service directly..."

build_scripts/install_comfy.sh

@@ -22,8 +22,7 @@ if [ -n "$ESD_COMMIT_ID" ]; then
 fi
 
 cp stable-diffusion-aws-extension/build_scripts/comfy/serve.py ComfyUI/
-cp stable-diffusion-aws-extension/build_scripts/comfy/comfy_sagemaker_proxy.py ComfyUI/custom_nodes/
+cp stable-diffusion-aws-extension/build_scripts/comfy/comfy_proxy.py ComfyUI/custom_nodes/
-cp stable-diffusion-aws-extension/build_scripts/comfy/comfy_local_proxy.py ComfyUI/custom_nodes/
 cp -R stable-diffusion-aws-extension/build_scripts/comfy/ComfyUI-AWS-Extension ComfyUI/custom_nodes/ComfyUI-AWS-Extension
 
 rm -rf stable-diffusion-aws-extension
@@ -40,7 +39,7 @@ git clone https://github.com/AIGODLIKE/AIGODLIKE-ComfyUI-Translation.git custom_
 git clone https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite.git custom_nodes/ComfyUI-VideoHelperSuite
 git clone https://github.com/Kosinkadink/ComfyUI-AnimateDiff-Evolved.git custom_nodes/ComfyUI-AnimateDiff-Evolved
 
-if [ "$ON_DOCKER" == "true" ]; then
+if [ "$ON_SAGEMAKER" == "true" ]; then
   python3 -m venv venv
   source venv/bin/activate
   pip install --upgrade pip

build_scripts/install_sd.sh

@@ -32,7 +32,7 @@ if [ -n "$ESD_COMMIT_ID" ]; then
 fi
 
 # remove unused files for docker layer reuse
-if [ "$ON_DOCKER" == "true" ]; then
+if [ "$ON_SAGEMAKER" == "true" ]; then
   rm -rf docs
   rm -rf infrastructure
   rm -rf middleware_api

comfy_start.sh

@@ -2,25 +2,26 @@
 
 set -euxo pipefail
 
-source /etc/environment
+if [ -f "/etc/environment" ]; then
+    source /etc/environment
+fi
 
-export CONTAINER_NAME='comfy_ec2'
+export CONTAINER_NAME='esd_comfy'
 export ACCOUNT_ID=$(aws sts get-caller-identity --query "Account" --output text)
 export AWS_REGION=$(aws configure get region)
 
-repository_name="comfy-ec2"
+image="$ACCOUNT_ID.dkr.ecr.$AWS_REGION.amazonaws.com/$CONTAINER_NAME:latest"
-image="$ACCOUNT_ID.dkr.ecr.$AWS_REGION.amazonaws.com/$repository_name:latest"
 
 docker stop "$CONTAINER_NAME" || true
 docker rm "$CONTAINER_NAME" || true
 
 # Check if the repository already exists
-if aws ecr describe-repositories --region "$AWS_REGION" --repository-names "$repository_name" >/dev/null 2>&1; then
+if aws ecr describe-repositories --region "$AWS_REGION" --repository-names "$CONTAINER_NAME" >/dev/null 2>&1; then
-    echo "ECR repository '$repository_name' already exists."
+    echo "ECR repository '$CONTAINER_NAME' already exists."
 else
-    echo "ECR repository '$repository_name' does not exist. Creating..."
+    echo "ECR repository '$CONTAINER_NAME' does not exist. Creating..."
-    aws ecr create-repository --repository-name --region "$AWS_REGION" "$repository_name"
+    aws ecr create-repository --repository-name --region "$AWS_REGION" "$CONTAINER_NAME" | jq .
-    echo "ECR repository '$repository_name' created successfully."
+    echo "ECR repository '$CONTAINER_NAME' created successfully."
 fi
 
 aws ecr get-login-password --region "$AWS_REGION" | docker login --username AWS --password-stdin "366590864501.dkr.ecr.$AWS_REGION.amazonaws.com"
@@ -32,7 +33,7 @@ docker build -f Dockerfile.comfy \
 image_hash=$(docker inspect "$image"  | jq -r ".[0].Id")
 image_hash=${image_hash:7}
 
-release_image="$ACCOUNT_ID.dkr.ecr.$AWS_REGION.amazonaws.com/$repository_name:$image_hash"
+release_image="$ACCOUNT_ID.dkr.ecr.$AWS_REGION.amazonaws.com/$CONTAINER_NAME:$image_hash"
 docker tag "$image" "$release_image"
 
 aws ecr get-login-password --region "$AWS_REGION" | docker login --username AWS --password-stdin "$ACCOUNT_ID.dkr.ecr.$AWS_REGION.amazonaws.com"
@@ -46,13 +47,16 @@ echo "Starting container..."
 local_volume="./ComfyUI"
 # local vol can be replace with your local directory
 
+#  -v ./build_scripts/comfy/comfy_local_proxy.py:/home/ubuntu/ComfyUI/custom_nodes/comfy_local_proxy.py \
+
 docker run -v ~/.aws:/root/.aws \
-           -v $local_volume:/home/ubuntu/ComfyUI \
+           -v "$local_volume":/home/ubuntu/ComfyUI \
+           -v ./build_scripts/inference/start.sh:/start.sh \
            --gpus all \
            -e "IMAGE_HASH=$release_image" \
            -e "ESD_VERSION=$ESD_VERSION" \
            -e "SERVICE_TYPE=comfy" \
-           -e "COMFY_EC2=true" \
+           -e "ON_EC2=true" \
            -e "S3_BUCKET_NAME=$COMFY_BUCKET_NAME" \
            -e "AWS_REGION=$AWS_REGION" \
            -e "AWS_DEFAULT_REGION=$AWS_REGION" \
@@ -61,6 +65,7 @@ docker run -v ~/.aws:/root/.aws \
            -e "COMFY_ENDPOINT=$COMFY_ENDPOINT" \
            -e "COMFY_BUCKET_NAME=$COMFY_BUCKET_NAME" \
            -e "PROCESS_NUMBER=$PROCESS_NUMBER" \
+           -e "WORKFLOW_NAME=$WORKFLOW_NAME" \
            --name "$CONTAINER_NAME" \
            -p 8188-8288:8188-8288 \
            "$image"

infrastructure/src/api/workflows/delete-workflows.ts

@@ -0,0 +1,170 @@
+import {PythonFunction} from '@aws-cdk/aws-lambda-python-alpha';
+import {Aws, aws_lambda, Duration} from 'aws-cdk-lib';
+import {JsonSchemaType, JsonSchemaVersion, LambdaIntegration, Model, Resource} from 'aws-cdk-lib/aws-apigateway';
+import {Table} from 'aws-cdk-lib/aws-dynamodb';
+import {Effect, PolicyStatement, Role, ServicePrincipal} from 'aws-cdk-lib/aws-iam';
+import {Architecture, LayerVersion, Runtime} from 'aws-cdk-lib/aws-lambda';
+import {Construct} from 'constructs';
+import {ApiModels} from '../../shared/models';
+import {SCHEMA_WORKFLOW_NAME} from '../../shared/schema';
+import {ApiValidators} from '../../shared/validator';
+
+export interface DeleteWorkflowsApiProps {
+    router: Resource;
+    httpMethod: string;
+    workflowsTable: Table;
+    multiUserTable: Table;
+    commonLayer: LayerVersion;
+}
+
+export class DeleteWorkflowsApi {
+    private readonly router: Resource;
+    private readonly httpMethod: string;
+    private readonly scope: Construct;
+    private readonly workflowsTable: Table;
+    private readonly multiUserTable: Table;
+    private readonly layer: LayerVersion;
+    private readonly baseId: string;
+
+    constructor(scope: Construct, id: string, props: DeleteWorkflowsApiProps) {
+        this.scope = scope;
+        this.baseId = id;
+        this.router = props.router;
+        this.httpMethod = props.httpMethod;
+        this.workflowsTable = props.workflowsTable;
+        this.multiUserTable = props.multiUserTable;
+        this.layer = props.commonLayer;
+
+        const lambdaFunction = this.apiLambda();
+
+        const lambdaIntegration = new LambdaIntegration(
+            lambdaFunction,
+            {
+                proxy: true,
+            },
+        );
+
+        this.router.addMethod(this.httpMethod, lambdaIntegration, {
+            apiKeyRequired: true,
+            requestValidator: ApiValidators.bodyValidator,
+            requestModels: {
+                'application/json': this.createRequestBodyModel(),
+            },
+            operationName: 'DeleteWorkflows',
+            methodResponses: [
+                ApiModels.methodResponses204(),
+                ApiModels.methodResponses400(),
+                ApiModels.methodResponses401(),
+                ApiModels.methodResponses403(),
+            ],
+        });
+    }
+
+    private iamRole(): Role {
+
+        const newRole = new Role(this.scope, `${this.baseId}-role`, {
+            assumedBy: new ServicePrincipal('lambda.amazonaws.com'),
+        });
+
+
+        newRole.addToPolicy(new PolicyStatement({
+            actions: [
+                'dynamodb:Query',
+                'dynamodb:GetItem',
+                'dynamodb:PutItem',
+                'dynamodb:DeleteItem',
+                'dynamodb:UpdateItem',
+                'dynamodb:Describe*',
+                'dynamodb:List*',
+            ],
+            resources: [
+                this.workflowsTable.tableArn,
+                `${this.workflowsTable.tableArn}/*`,
+                this.multiUserTable.tableArn,
+            ],
+        }));
+
+        newRole.addToPolicy(new PolicyStatement({
+            actions: [
+                's3:Get*',
+                's3:List*',
+                's3:PutObject',
+                's3:GetObject',
+                's3:DeleteObject',
+            ],
+            resources: [
+                '*',
+            ],
+        }));
+
+        newRole.addToPolicy(new PolicyStatement({
+            effect: Effect.ALLOW,
+            actions: [
+                'cloudwatch:DeleteAlarms',
+                'cloudwatch:DescribeAlarms',
+                'cloudwatch:DeleteDashboards',
+            ],
+            resources: [
+                '*',
+            ],
+        }));
+
+        newRole.addToPolicy(new PolicyStatement({
+            effect: Effect.ALLOW,
+            actions: [
+                'logs:CreateLogGroup',
+                'logs:CreateLogStream',
+                'logs:PutLogEvents',
+                'logs:DeleteLogGroup',
+            ],
+            resources: [`arn:${Aws.PARTITION}:logs:${Aws.REGION}:${Aws.ACCOUNT_ID}:log-group:*:*`],
+        }));
+
+        return newRole;
+    }
+
+    private createRequestBodyModel(): Model {
+        return new Model(this.scope, `${this.baseId}-model`, {
+            restApi: this.router.api,
+            modelName: this.baseId,
+            description: `Request Model ${this.baseId}`,
+            schema: {
+                schema: JsonSchemaVersion.DRAFT7,
+                title: this.baseId,
+                type: JsonSchemaType.OBJECT,
+                properties: {
+                    workflow_name_list: {
+                        type: JsonSchemaType.ARRAY,
+                        items: SCHEMA_WORKFLOW_NAME,
+                        minItems: 1,
+                        maxItems: 10,
+                    },
+                },
+                required: [
+                    'workflow_name_list',
+                ],
+            },
+            contentType: 'application/json',
+        });
+    }
+
+    private apiLambda() {
+        return new PythonFunction(this.scope, `${this.baseId}-lambda`, {
+            entry: '../middleware_api/workflows',
+            architecture: Architecture.X86_64,
+            runtime: Runtime.PYTHON_3_10,
+            index: 'delete_workflows.py',
+            handler: 'handler',
+            timeout: Duration.seconds(900),
+            role: this.iamRole(),
+            memorySize: 2048,
+            tracing: aws_lambda.Tracing.ACTIVE,
+            layers: [this.layer],
+            environment:{
+                WORKFLOWS_TABLE: this.workflowsTable.tableName,
+            }
+        });
+    }
+
+
+}

infrastructure/src/shared/workflow.ts

@@ -6,6 +6,7 @@ import {Construct} from 'constructs';
 import {ResourceProvider} from './resource-provider';
 import {CreateWorkflowApi} from "../api/workflows/create-workflow";
 import {ListWorkflowsApi} from "../api/workflows/list-workflows";
+import {DeleteWorkflowsApi} from "../api/workflows/delete-workflows";
 
 export interface WorkflowProps extends StackProps {
     routers: { [key: string]: Resource };
@@ -45,6 +46,17 @@ export class Workflow {
             },
         );
 
+        new DeleteWorkflowsApi(
+            scope, 'DeleteWorkflows',
+            {
+                workflowsTable: props.workflowsTable,
+                commonLayer: props.commonLayer,
+                multiUserTable: props.multiUserTable,
+                httpMethod: 'DELETE',
+                router: props.routers.workflows,
+            },
+        );
+
     }
 
 }

middleware_api/endpoints/create_endpoint.py

@@ -245,13 +245,13 @@ def _create_sagemaker_model(name, model_data_url, endpoint_name, endpoint_id, ev
         'ESD_VERSION': esd_version,
         'ESD_COMMIT_ID': esd_commit_id,
         'SERVICE_TYPE': event.service_type,
-        'ON_DOCKER': 'true',
+        'ON_SAGEMAKER': 'true',
         'AWS_REGION': aws_region,
         'AWS_DEFAULT_REGION': aws_region,
     }
 
     if event.workflow:
-        environment['APP_SOURCE'] = event.workflow.s3_location
+        environment['WORKFLOW_NAME'] = event.workflow.name
         environment['APP_CWD'] = '/home/ubuntu/ComfyUI'
 
     primary_container = {

middleware_api/service/oas.py

@@ -485,6 +485,11 @@ operations = {
         tags=["Workflows"],
         description="List Workflows with Parameters",
     ),
+    "DeleteWorkflows": APISchema(
+        summary="Delete Workflows",
+        tags=["Workflows"],
+        description="Delete specify Workflows",
+    ),
 }
 
 

middleware_api/workflows/delete_workflows.py

@@ -0,0 +1,47 @@
+import json
+import logging
+import os
+from dataclasses import dataclass
+
+import boto3
+from aws_lambda_powertools import Tracer
+
+from common.ddb_service.client import DynamoDbUtilsService
+from common.response import no_content
+from libs.utils import response_error
+
+tracer = Tracer()
+workflows_table = os.environ.get('WORKFLOWS_TABLE')
+
+logger = logging.getLogger(__name__)
+logger.setLevel(os.environ.get('LOG_LEVEL') or logging.ERROR)
+
+ddb_service = DynamoDbUtilsService(logger=logger)
+esd_version = os.environ.get("ESD_VERSION")
+s3_resource = boto3.resource('s3')
+bucket_name = os.environ.get('S3_BUCKET_NAME')
+s3_bucket = s3_resource.Bucket(bucket_name)
+
+
+@dataclass
+class DeleteWorkflowsEvent:
+    workflow_name_list: [str]
+
+
+@tracer.capture_lambda_handler
+def handler(raw_event, ctx):
+    try:
+        logger.info(json.dumps(raw_event))
+
+        event = DeleteWorkflowsEvent(**json.loads(raw_event['body']))
+
+        for name in event.workflow_name_list:
+            s3_bucket.objects.filter(Prefix=f"comfy/workflows/{name}/").delete()
+            ddb_service.delete_item(
+                table=workflows_table,
+                keys={'name': name},
+            )
+
+        return no_content(message="Workflows Deleted")
+    except Exception as e:
+        return response_error(e)

scripts/api.py

@@ -526,8 +526,8 @@ try:
     import modules.script_callbacks as script_callbacks
 
     script_callbacks.on_app_started(sagemaker_api)
-    on_docker = os.environ.get('ON_DOCKER', "false")
+    on_sagemaker = os.environ.get('ON_SAGEMAKER', "false")
-    if on_docker == "true":
+    if on_sagemaker == "true":
         from modules import shared
         shared.opts.data.update(control_net_max_models_num=10)
         script_callbacks.on_app_started(move_model_to_tmp)

scripts/main.py

@@ -855,8 +855,8 @@ class SageMakerUI(scripts.Script):
         return sagemaker_inputs_components
 
     def before_process(self, p, *args):
-        on_docker = os.environ.get('ON_DOCKER', "false")
+        on_sagemaker = os.environ.get('ON_SAGEMAKER', "false")
-        if on_docker == "true":
+        if on_sagemaker == "true":
             return
 
         # check if endpoint is InService
@@ -1055,7 +1055,7 @@ def fetch_user_data():
         time.sleep(30)
 
 
-if os.environ.get('ON_DOCKER', "false") != "true":
+if os.environ.get('ON_SAGEMAKER', "false") != "true":
     from aws_extension.auth_service.simple_cloud_auth import cloud_auth_manager
     if cloud_auth_manager.enableAuth:
         cmd_opts.gradio_auth = cloud_auth_manager.create_config()

test/test_02_api_base/test_12_workflows.py

@@ -21,26 +21,45 @@ class TestComfyWorkflowApiBase:
         resp = self.api.create_workflow()
         assert resp.status_code == 403, resp.dumps()
 
-    def test_2_list_workflows_without_key(self):
+    def test_2_create_workflow_with_bad_key(self):
-        resp = self.api.list_workflows()
-        assert resp.status_code == 403, resp.dumps()
-
-    def test_3_create_workflow_with_bad_key(self):
         headers = {'x-api-key': "bad_key"}
         resp = self.api.create_workflow(headers)
         assert resp.status_code == 403, resp.dumps()
 
-    def test_4_list_workflows_with_bad_key(self):
+    def test_3_create_workflow_with_bad_request(self):
-        headers = {'x-api-key': "bad_key"}
-        resp = self.api.list_workflows(headers)
-        assert resp.status_code == 403, resp.dumps()
-
-    def test_5_create_workflow_with_bad_request(self):
         headers = {'x-api-key': config.api_key}
         resp = self.api.create_workflow(headers)
         assert resp.status_code == 400, resp.dumps()
 
-    def test_6_list_executes_with_ok(self):
+    def test_4_list_workflows_without_key(self):
+        resp = self.api.list_workflows()
+        assert resp.status_code == 403, resp.dumps()
+
+    def test_5_list_workflows_with_bad_key(self):
+        headers = {'x-api-key': "bad_key"}
+        resp = self.api.list_workflows(headers)
+        assert resp.status_code == 403, resp.dumps()
+
+    def test_6_list_workflows_with_ok(self):
         headers = {'x-api-key': config.api_key}
         resp = self.api.list_workflows(headers)
         assert resp.status_code == 200, resp.dumps()
+
+    def test_7_delete_workflows_without_key(self):
+        resp = self.api.delete_workflows()
+        assert resp.status_code == 403, resp.dumps()
+
+    def test_8_delete_workflows_with_bad_key(self):
+        headers = {'x-api-key': "bad_key"}
+        resp = self.api.delete_workflows(headers)
+        assert resp.status_code == 403, resp.dumps()
+
+    def test_9_delete_workflows_with_ok(self):
+        headers = {'x-api-key': config.api_key}
+        data = {
+            "workflow_name_list": [
+                "workflow_name"
+            ],
+        }
+        resp = self.api.delete_workflows(headers=headers, data=data)
+        assert resp.status_code == 204, resp.dumps()

test/utils/api.py

@@ -92,6 +92,15 @@ class Api:
             data=data
         )
 
+    def delete_workflows(self, headers=None, data=None):
+        return self.req(
+            "DELETE",
+            "workflows",
+            headers=headers,
+            operation_id='DeleteWorkflows',
+            data=data
+        )
+
     def delete_users(self, headers=None, data=None):
         return self.req(
             "DELETE",

workshop/Dockerfile.comfy

@@ -1,7 +0,0 @@
-ARG AWS_REGION
-FROM 366590864501.dkr.ecr.$AWS_REGION.amazonaws.com/esd-inference:dev
-
-# TODO BYOC
-RUN apt-get update -y && \
-    apt-get install ffmpeg -y && \
-    rm -rf /var/lib/apt/lists/*

workshop/comfy.yaml

@@ -56,6 +56,9 @@ Parameters:
       - latest
       - dev
     Default: latest
+  WorkflowName:
+    Description: Bind Workflow Name
+    Type: String
 
 Mappings:
   RegionToAmiId:
@@ -218,6 +221,7 @@ Resources:
           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
           
@@ -243,7 +247,7 @@ Resources:
           StartLimitIntervalSec=0
 
           [Service]
-          WorkingDirectory=/root/stable-diffusion-aws-extension/workshop/
+          WorkingDirectory=/root/stable-diffusion-aws-extension/
           ExecStart=bash comfy_start.sh
           Type=simple
           Restart=always