aws-sdk-go経由で AWS bedrock / Stable Diffusionで画像生成するメモ

個人用メモ

• 上位モデルは現在予約済のスループットでしかサポートされていない?
• goに関してはbedrock でなく bedrockruntime パッケージを使う。
• 画像はbase64データとしてSDKより直接返ってくる
  • APIリクエストのオプションはstability.ai と同様

package main

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/bedrockruntime"
	"github.com/wano/contextlog/clog"
	"os"
	"path/filepath"
	"time"
)

func main() {
	sess, _ := session.NewSession(&aws.Config{
		Region: aws.String("us-east-1")},
	)

	brr := bedrockruntime.New(sess)

	type Prompt struct {
		Text string `json:"text"`
	}

	type ReqBody struct {
		TextPrompts []Prompt `json:"text_prompts"`
		CfgScale    int      `json:"cfg_scale"`
		Seed        int      `json:"seed"`
		Steps       int      `json:"steps"`
		Width       int      `json:"width"`
		Height      int      `json:"height"`
	}

	reqBody := ReqBody{
		TextPrompts: []Prompt{
			{
				Text: " Music Album Jacket for HIPHOP ",
			},
		},
		CfgScale: 30,
		Seed:     749500536,
		Steps:    130,
		Width:    512,
		Height:   512,
	}

	rb, _ := json.Marshal(reqBody)

	resp, err := brr.InvokeModel(&bedrockruntime.InvokeModelInput{
		ModelId:     aws.String(`stability.stable-diffusion-xl-v0`),
		ContentType: aws.String(`application/json`),
		Accept:      aws.String(`application/json`),
		//Body:        []byte("{\"text_prompts\":[{\"text\":\" Music Album Jacket for HIPHOP \"}],\"cfg_scale\":30,\"seed\":749500524,\"steps\":130}"),
		Body: rb,
	})
	if err != nil {
		panic(err)
	}

	type Artifact struct {
		Seed   string `json:"seed"`
		Base64 string `json:"base64"`
	}

	type Result struct {
		Result    string     `json:"result"`
		Artifacts []Artifact `json:"artifacts"`
	}

	rr := Result{}
	_ = json.Unmarshal(resp.Body, &rr)

	if rr.Result != `success` {
		panic(rr.Result)
	}

	b64 := rr.Artifacts[0].Base64

	// Base64データをデコードする
	decodedData, err := base64.StdEncoding.DecodeString(b64)
	if err != nil {
		clog.Fatal("Error decoding image:", err)
	}

	os.MkdirAll("./output", 0777)

	fileP := filepath.Join("./output", fmt.Sprintf("%s.jpg", time.Now().Format("20060102150405")))

	// デコードしたデータをJPEGファイルとして保存する
	err = os.WriteFile(fileP, decodedData, 0777)
	if err != nil {
		clog.Fatal("Error writing to file:", err)
	}

}