Skip to content

Commit 1c4a6f7

Browse files
philipwsphilipws
committed
Initial Stability Diffusion v3 support
1 parent b81a67b commit 1c4a6f7

File tree

3 files changed

+269
-29
lines changed

3 files changed

+269
-29
lines changed

‎infrastructure/lib/features/readable/generate.ts‎

Lines changed: 29 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import { dt_readableWorkflow as dt_readableWorkflow_amazon_titanText } from "./v
2020
import { dt_readableWorkflow as dt_readableWorkflow_anthropic_claudeText } from "./vendor/text.anthropic.claude-v2";
2121
import { dt_readableWorkflow as dt_readableWorkflow_anthropic_claude3Text } from "./vendor/text.anthropic.claude-3-*-*-v1";
2222
import { dt_readableWorkflow as dt_readableWorkflow_stabilityai_stableDiffusion } from "./vendor/image.stability.stable-diffusion-xl-v1";
23+
import { dt_readableWorkflow as dt_readableWorkflow_stabilityai_stableDiffusion_3 } from "./vendor/image.stability.sd3-*";
2324

2425
export interface props {
2526
bedrockRegion: string;
@@ -178,6 +179,17 @@ export class dt_readableWorkflowGenerate extends Construct {
178179
removalPolicy: props.removalPolicy,
179180
},
180181
);
182+
// MODEL WORKFLOWS | IMAGE | STABILITYAI v3
183+
const workflow_stabilityai_3 =
184+
new dt_readableWorkflow_stabilityai_stableDiffusion_3(
185+
this,
186+
"workflow_stabilityai_3",
187+
{
188+
bedrockRegion: props.bedrockRegion,
189+
contentBucket: props.contentBucket,
190+
removalPolicy: props.removalPolicy,
191+
},
192+
);
181193

182194
//
183195
// STATE MACHINE
@@ -209,44 +221,35 @@ export class dt_readableWorkflowGenerate extends Construct {
209221
workflow_stabilityai.modelChoiceCondition,
210222
workflow_stabilityai.invokeModel,
211223
)
224+
.when(
225+
workflow_stabilityai_3.modelChoiceCondition,
226+
workflow_stabilityai_3.invokeModel,
227+
)
212228
.otherwise(failUnrecognisedModel),
213229
},
214230
).StateMachine;
215231

232+
const workflows = [
233+
workflow_amazon_text,
234+
workflow_amazon_image,
235+
workflow_anthropic,
236+
workflow_anthropic3,
237+
workflow_stabilityai,
238+
workflow_stabilityai_3
239+
];
216240
NagSuppressions.addResourceSuppressions(
217241
this.sfnMain,
218242
[
219243
{
220244
id: "AwsSolutions-IAM5",
221-
reason:
222-
"Permission scoped to project specific resources. Execution ID unknown at deploy time.",
223-
appliesTo: [
224-
`Resource::arn:<AWS::Partition>:states:<AWS::Region>:<AWS::AccountId>:execution:{"Fn::Select":[6,{"Fn::Split":[":",{"Ref":"${cdk.Stack.of(
225-
this,
226-
).getLogicalId(
227-
workflow_amazon_text.sfnMain.node.defaultChild as cdk.CfnElement,
228-
)}"}]}]}*`,
229-
`Resource::arn:<AWS::Partition>:states:<AWS::Region>:<AWS::AccountId>:execution:{"Fn::Select":[6,{"Fn::Split":[":",{"Ref":"${cdk.Stack.of(
230-
this,
231-
).getLogicalId(
232-
workflow_amazon_image.sfnMain.node.defaultChild as cdk.CfnElement,
233-
)}"}]}]}*`,
234-
`Resource::arn:<AWS::Partition>:states:<AWS::Region>:<AWS::AccountId>:execution:{"Fn::Select":[6,{"Fn::Split":[":",{"Ref":"${cdk.Stack.of(
235-
this,
236-
).getLogicalId(
237-
workflow_anthropic.sfnMain.node.defaultChild as cdk.CfnElement,
238-
)}"}]}]}*`,
239-
`Resource::arn:<AWS::Partition>:states:<AWS::Region>:<AWS::AccountId>:execution:{"Fn::Select":[6,{"Fn::Split":[":",{"Ref":"${cdk.Stack.of(
240-
this,
241-
).getLogicalId(
242-
workflow_anthropic3.sfnMain.node.defaultChild as cdk.CfnElement,
243-
)}"}]}]}*`,
245+
reason: "Permission scoped to project specific resources. Execution ID unknown at deploy time.",
246+
appliesTo: workflows.map(workflow =>
244247
`Resource::arn:<AWS::Partition>:states:<AWS::Region>:<AWS::AccountId>:execution:{"Fn::Select":[6,{"Fn::Split":[":",{"Ref":"${cdk.Stack.of(
245248
this,
246249
).getLogicalId(
247-
workflow_stabilityai.sfnMain.node.defaultChild as cdk.CfnElement,
248-
)}"}]}]}*`,
249-
],
250+
workflow.sfnMain.node.defaultChild as cdk.CfnElement,
251+
)}"}]}]}*`
252+
)
250253
},
251254
],
252255
true,

‎infrastructure/lib/features/readable/vendor/image.stability.sd3-*.ts‎

Lines changed: 235 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,235 @@
1+
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: MIT-0
3+
4+
import { Construct } from "constructs";
5+
import * as cdk from "aws-cdk-lib";
6+
import { NagSuppressions } from "cdk-nag";
7+
8+
import {
9+
aws_stepfunctions as sfn,
10+
aws_stepfunctions_tasks as tasks,
11+
aws_iam as iam,
12+
aws_s3 as s3,
13+
} from "aws-cdk-lib";
14+
15+
import { Bucket } from "../enum";
16+
import { dt_lambda } from "../../../components/lambda";
17+
import { dt_stepfunction } from "../../../components/stepfunction";
18+
19+
export interface props {
20+
bedrockRegion: string;
21+
contentBucket: s3.Bucket;
22+
removalPolicy: cdk.RemovalPolicy;
23+
}
24+
25+
enum Strings {
26+
modelVendor = "stability",
27+
modelNamePrefix = "sd3-",
28+
}
29+
30+
const id = "v3";
31+
32+
export class dt_readableWorkflow extends Construct {
33+
public readonly invokeModel: tasks.StepFunctionsStartExecution;
34+
public readonly modelChoiceCondition: sfn.Condition;
35+
public readonly sfnMain: sfn.StateMachine;
36+
37+
constructor(scope: Construct, id: string, props: props) {
38+
super(scope, id);
39+
40+
// LAMBDA
41+
// LAMBDA | INVOKE BEDROCK
42+
// LAMBDA | INVOKE BEDROCK | ROLE
43+
const invokeBedrockLambdaRole = new iam.Role(
44+
this,
45+
"invokeBedrockSaveToS3",
46+
{
47+
// ASM-L6 // ASM-L8
48+
assumedBy: new iam.ServicePrincipal("lambda.amazonaws.com"),
49+
description: "Lambda Role (Invoke Bedrock API & save file to S3)",
50+
},
51+
);
52+
53+
// LAMBDA | INVOKE BEDROCK | POLICY
54+
const permitInvokeBedrockModel = new iam.Policy(
55+
this,
56+
"permitInvokeModels",
57+
{
58+
policyName: `Invoke-${Strings.modelVendor}-${Strings.modelNamePrefix}-models`,
59+
statements: [
60+
new iam.PolicyStatement({
61+
// ASM-IAM
62+
actions: ["bedrock:InvokeModel"],
63+
resources: [
64+
`arn:aws:bedrock:${props.bedrockRegion}::foundation-model/${Strings.modelVendor}.${Strings.modelNamePrefix}*`, // Foundational Models
65+
],
66+
}),
67+
],
68+
},
69+
);
70+
invokeBedrockLambdaRole.attachInlinePolicy(permitInvokeBedrockModel);
71+
NagSuppressions.addResourceSuppressions(
72+
permitInvokeBedrockModel,
73+
[
74+
{
75+
id: "AwsSolutions-IAM5",
76+
reason: "Preferred model for prompt is unknown at deploy time",
77+
appliesTo: [
78+
`Resource::arn:aws:bedrock:${props.bedrockRegion}::foundation-model/${Strings.modelVendor}.${Strings.modelNamePrefix}*`,
79+
],
80+
},
81+
],
82+
true,
83+
);
84+
const permitPutS3 = new iam.Policy(this, "permitPutS3", {
85+
policyName: "S3-PutObject",
86+
statements: [
87+
new iam.PolicyStatement({
88+
// ASM-IAM
89+
actions: ["s3:PutObject"],
90+
resources: [
91+
`${props.contentBucket.bucketArn}/${Bucket.PREFIX_PRIVATE}/*`,
92+
],
93+
}),
94+
],
95+
});
96+
invokeBedrockLambdaRole.attachInlinePolicy(permitPutS3);
97+
NagSuppressions.addResourceSuppressions(
98+
permitPutS3,
99+
[
100+
{
101+
id: "AwsSolutions-IAM5",
102+
reason:
103+
"Scoped to project specific resources. Unknown filenames at deploy time.",
104+
appliesTo: [
105+
"Resource::<basereadablecontentBucketbucket6155EB26.Arn>/private/*",
106+
],
107+
},
108+
],
109+
true,
110+
);
111+
112+
// LAMBDA | INVOKE BEDROCK | FUNCTION
113+
const invokeBedrockLambda = new dt_lambda(this, "invokeBedrockLambda", {
114+
role: invokeBedrockLambdaRole,
115+
path: "lambda/invokeBedrockSaveToS3",
116+
description: "Invoke Bedrock API & save file to S3",
117+
environment: {
118+
BEDROCK_REGION: props.bedrockRegion,
119+
},
120+
bundlingNodeModules: [
121+
"@aws-sdk/client-bedrock-runtime",
122+
"@aws-sdk/client-s3",
123+
],
124+
timeout: cdk.Duration.seconds(60),
125+
});
126+
127+
//
128+
// STATE MACHINE
129+
// STATE MACHINE | TASKS
130+
// STATE MACHINE | TASKS | createPrompt
131+
const createPrompt = new sfn.Pass(this, "createPrompt", {
132+
resultPath: "$.prompt",
133+
parameters: {
134+
Payload: {
135+
prompt: sfn.JsonPath.stringAt("$.jobDetails.prePrompt"),
136+
},
137+
},
138+
});
139+
// STATE MACHINE | TASKS | createBody
140+
const createBody = new sfn.Pass(this, "createBody", {
141+
resultPath: "$.body",
142+
parameters: {
143+
Payload: sfn.JsonPath.jsonMerge(
144+
sfn.JsonPath.objectAt("$.prompt.Payload"),
145+
sfn.JsonPath.objectAt("$.jobDetails.parameters"),
146+
),
147+
},
148+
});
149+
150+
// STATE MACHINE | TASKS | invokeBedrock
151+
const invokeBedrock = new tasks.LambdaInvoke(this, "invokeBedrock", {
152+
lambdaFunction: invokeBedrockLambda.lambdaFunction,
153+
resultPath: "$.invokeBedrock",
154+
resultSelector: {
155+
"Payload.$": "$.Payload",
156+
},
157+
payload: sfn.TaskInput.fromObject({
158+
ModelId: sfn.JsonPath.objectAt("$.jobDetails.modelId"),
159+
Body: sfn.JsonPath.stringAt("$.body.Payload"),
160+
PathToResult: "images.0",
161+
ResultS3Bucket: props.contentBucket.bucketName,
162+
ResultS3Key: sfn.JsonPath.format(
163+
`${Bucket.PREFIX_PRIVATE}/{}/{}/{}_{}.png`,
164+
sfn.JsonPath.stringAt("$.jobDetails.identity"),
165+
sfn.JsonPath.stringAt("$.jobDetails.id"),
166+
sfn.JsonPath.stringAt("$.jobDetails.itemId"),
167+
sfn.JsonPath.stringAt("$$.Execution.StartTime"),
168+
),
169+
ItemId: sfn.JsonPath.stringAt("$.jobDetails.itemId"),
170+
}),
171+
});
172+
173+
// STATE MACHINE | TASKS | filterOutput
174+
const filterOutput = new sfn.Pass(this, "filterOutput", {
175+
parameters: {
176+
payload: sfn.JsonPath.stringAt("$.invokeBedrock.Payload.key"),
177+
},
178+
});
179+
180+
// STATE MACHINE | DEF
181+
this.sfnMain = new dt_stepfunction(
182+
this,
183+
`${cdk.Stack.of(this).stackName}_Readable_${Strings.modelVendor}_${id}`,
184+
{
185+
nameSuffix: `Readable_${Strings.modelVendor}_${id}`,
186+
removalPolicy: props.removalPolicy,
187+
definition: createPrompt
188+
.next(createBody)
189+
.next(invokeBedrock)
190+
.next(filterOutput),
191+
},
192+
).StateMachine;
193+
NagSuppressions.addResourceSuppressions(
194+
this.sfnMain,
195+
[
196+
{
197+
id: "AwsSolutions-IAM5",
198+
reason: "Permissions scoped to dedicated resources.",
199+
appliesTo: [
200+
`Resource::<${cdk.Stack.of(this).getLogicalId(
201+
invokeBedrockLambda.lambdaFunction.node
202+
.defaultChild as cdk.CfnElement,
203+
)}.Arn>:*`,
204+
],
205+
},
206+
],
207+
true,
208+
);
209+
210+
// PARENT
211+
// PARENT | CHOICE FILTER
212+
this.modelChoiceCondition = sfn.Condition.stringMatches(
213+
"$.jobDetails.modelId",
214+
`${Strings.modelVendor}.${Strings.modelNamePrefix}*`,
215+
);
216+
// PARENT | TASK
217+
this.invokeModel = new tasks.StepFunctionsStartExecution(
218+
this,
219+
`invokeModel_${Strings.modelVendor}_${id}`,
220+
{
221+
stateMachine: this.sfnMain,
222+
resultSelector: {
223+
"Payload.$": "$.Output.payload",
224+
},
225+
resultPath: "$.invokeModel",
226+
integrationPattern: sfn.IntegrationPattern.RUN_JOB,
227+
input: sfn.TaskInput.fromObject({
228+
jobDetails: sfn.JsonPath.objectAt("$.jobDetails"),
229+
}),
230+
},
231+
);
232+
233+
// END
234+
}
235+
}

‎infrastructure/lib/features/readable/vendor/image.stability.stable-diffusion-xl-v1.ts‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@ enum Strings {
2727
modelNamePrefix = "stable-diffusion-",
2828
}
2929

30+
const id = "v1";
31+
3032
export class dt_readableWorkflow extends Construct {
3133
public readonly invokeModel: tasks.StepFunctionsStartExecution;
3234
public readonly modelChoiceCondition: sfn.Condition;
@@ -51,7 +53,7 @@ export class dt_readableWorkflow extends Construct {
5153
// LAMBDA | INVOKE BEDROCK | POLICY
5254
const permitInvokeBedrockModel = new iam.Policy(
5355
this,
54-
"permitSfnSendSuccess",
56+
"permitInvokeModels",
5557
{
5658
policyName: `Invoke-${Strings.modelVendor}-${Strings.modelNamePrefix}-models`,
5759
statements: [
@@ -183,9 +185,9 @@ export class dt_readableWorkflow extends Construct {
183185
// STATE MACHINE | DEF
184186
this.sfnMain = new dt_stepfunction(
185187
this,
186-
`${cdk.Stack.of(this).stackName}_Readable_${Strings.modelVendor}`,
188+
`${cdk.Stack.of(this).stackName}_Readable_${Strings.modelVendor}_${id}`,
187189
{
188-
nameSuffix: `Readable_${Strings.modelVendor}`,
190+
nameSuffix: `Readable_${Strings.modelVendor}_${id}`,
189191
removalPolicy: props.removalPolicy,
190192
definition: createPrompt
191193
.next(createBody)

0 commit comments

Comments
 (0)