import boto3
import json
import base64
from pathlib import Path

# 创建 Bedrock Runtime 客户端，使用本机 AWS CLI 配置，指定 us-east-1 区域
bedrock_runtime = boto3.client(
    service_name='bedrock-runtime',
    region_name='us-east-1'
)

# 定义提示词
prompt1cn = "4000米海拔的青藏高原上的公路，公路是黑色沥青没有积雪，两边的地面是积雪覆盖。背后是七千米7000米以上的雪山被太阳照亮，但是途中画面角度不直接看到太阳"
prompt2cn = "日照金山，画面地点在雪域高原，背后是8000米的雪山尖峰，金色的太阳光照亮了顶部"
prompt1en = "A black asphalt highway on the Tibetan Plateau at 4000 meters altitude, the road is clear without snow, both sides covered with snow. Behind are majestic snow-capped mountains over 7000 meters high illuminated by sunlight, but the sun is not directly visible in the frame. Photorealistic, high quality, dramatic landscape."
prompt2en = "Golden hour sunlight on mountain peaks, majestic 8000-meter snow-capped summits on the snowy plateau, golden sunlight illuminating the mountain tops, dramatic landscape photography"

prompt_submit = prompt1cn

# 准备请求体
request_body = {
    "taskType": "TEXT_IMAGE",
    "textToImageParams": {
        "text": prompt_submit
    },
    "imageGenerationConfig": {
        "numberOfImages": 1,
        "quality": "standard",
        "height": 1024,
        "width": 1024,
        "cfgScale": 8.0,
        "seed": 0
    }
}

# 调用 Nova Canvas 模型
model_id = "amazon.nova-canvas-v1:0"

print("正在调用 Nova Canvas 模型生成图片...")
print(f"提示词: {prompt_submit}")

response = bedrock_runtime.invoke_model(
    modelId=model_id,
    body=json.dumps(request_body)
)

# 解析响应
response_body = json.loads(response['body'].read())

# 获取生成的图片（base64 编码）
image_base64 = response_body['images'][0]

# 解码并保存图片
image_data = base64.b64decode(image_base64)
output_path = Path("output.png")

with open(output_path, "wb") as f:
    f.write(image_data)

print(f"图片已成功生成并保存到: {output_path.absolute()}")

import boto3
import json
import base64
from pathlib import Path

# 创建 Bedrock Runtime 客户端，使用本机 AWS CLI 配置，指定 us-west-2 区域
bedrock_runtime = boto3.client(
    service_name='bedrock-runtime',
    region_name='us-west-2'
)

# 定义提示词（使用英文以避免内容过滤）
prompt1 = "A black asphalt highway on the Tibetan Plateau at 4000 meters altitude, the road is clear without snow, both sides covered with snow. Behind are majestic snow-capped mountains over 7000 meters high illuminated by sunlight, but the sun is not directly visible in the frame. Photorealistic, high quality, dramatic landscape."
prompt2 = "Golden hour sunlight on mountain peaks, majestic 8000-meter snow-capped summits on the snowy plateau, golden sunlight illuminating the mountain tops, dramatic landscape photography"
prompt_submit = prompt1

# 准备请求体 - Stable Diffusion 3.5 Large 格式
request_body = {
    "prompt": prompt_submit,
    "mode": "text-to-image",
    "aspect_ratio": "1:1",
    "output_format": "png"
}

# 调用 Stable Diffusion 3.5 Large 模型
model_id = "stability.sd3-5-large-v1:0"

print("正在调用 Stable Diffusion 3.5 Large 模型生成图片...")
print(f"提示词: {prompt_submit}")

response = bedrock_runtime.invoke_model(
    modelId=model_id,
    body=json.dumps(request_body)
)

# 解析响应
response_body = json.loads(response['body'].read())

# 检查是否被内容过滤器拦截
if 'finish_reasons' in response_body and response_body['finish_reasons']:
    finish_reason = response_body['finish_reasons'][0]
    if finish_reason and ('Filter' in str(finish_reason) or 'CONTENT_FILTERED' in str(finish_reason)):
        print(f"错误: 请求被内容过滤器拦截 - {finish_reason}")
        print("建议: 尝试修改提示词或使用英文提示词")
        exit(1)

# 获取生成的图片（base64 编码）
# Stable Diffusion 3.5 的响应格式可能是 'image' 而不是 'images'
if 'images' in response_body:
    image_base64 = response_body['images'][0]
elif 'image' in response_body:
    image_base64 = response_body['image']
elif 'artifacts' in response_body:
    image_base64 = response_body['artifacts'][0]['base64']
else:
    print("响应内容:", json.dumps(response_body, indent=2, ensure_ascii=False))
    raise KeyError(f"无法在响应中找到图片数据。响应键: {list(response_body.keys())}")

# 解码并保存图片
image_data = base64.b64decode(image_base64)
output_path = Path("output_sd.png")

with open(output_path, "wb") as f:
    f.write(image_data)

print(f"图片已成功生成并保存到: {output_path.absolute()}")