Created
June 12, 2025 23:20
-
-
Save aphexlog/61cdb81038fda3d89a6b83460236a568 to your computer and use it in GitHub Desktop.
A lightweight Python script that uses AWS Bedrock to generate images from text prompts.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| #!/usr/bin/env python3 | |
| """ | |
| Simple image generator using AWS Bedrock | |
| """ | |
| import argparse | |
| import base64 | |
| import io | |
| import json | |
| import os | |
| from datetime import datetime | |
| import boto3 | |
| from PIL import Image | |
| def generate_image( | |
| prompt, model_id="amazon.titan-image-generator-v2:0", output_dir="generated_images" | |
| ): | |
| """Generate an image using AWS Bedrock based on the provided text prompt""" | |
| # Create Bedrock client | |
| bedrock_runtime = boto3.client( | |
| service_name="bedrock-runtime", | |
| region_name="us-east-1", # Change to your preferred region | |
| ) | |
| # Prepare request body based on the model | |
| if "stability" in model_id: | |
| request_body = { | |
| "text_prompts": [{"text": prompt}], | |
| "cfg_scale": 7, | |
| "steps": 30, | |
| "seed": 0, | |
| "width": 1024, | |
| "height": 1024, | |
| } | |
| elif "titan-image" in model_id: | |
| # Titan Image Generator format | |
| request_body = { | |
| "taskType": "TEXT_IMAGE", | |
| "textToImageParams": { | |
| "text": prompt, | |
| "negativeText": "blurry, bad quality, distorted", | |
| }, | |
| "imageGenerationConfig": { | |
| "numberOfImages": 1, | |
| "height": 1024, | |
| "width": 1024, | |
| "cfgScale": 8.0, | |
| }, | |
| } | |
| elif "claude" in model_id: | |
| # Claude format for image generation | |
| request_body = { | |
| "anthropic_version": "bedrock-2023-05-31", | |
| "max_tokens": 1024, | |
| "messages": [ | |
| { | |
| "role": "user", | |
| "content": [ | |
| {"type": "text", "text": f"Generate an image of: {prompt}"} | |
| ], | |
| } | |
| ], | |
| } | |
| else: | |
| # Generic format - try with simple prompt | |
| request_body = { | |
| "prompt": prompt, | |
| "negative_prompt": "blurry, bad quality, distorted", | |
| "image_size": "1024x1024", | |
| "num_images": 1, | |
| } | |
| # Invoke the model | |
| try: | |
| print(f"Attempting to generate image with model: {model_id}") | |
| print(f"Request body: {json.dumps(request_body, indent=2)}") | |
| response = bedrock_runtime.invoke_model( | |
| modelId=model_id, body=json.dumps(request_body) | |
| ) | |
| response_body = json.loads(response.get("body").read()) | |
| print(f"Response keys: {list(response_body.keys())}") | |
| # Handle response based on the model | |
| if "stability" in model_id: | |
| image_base64 = response_body.get("artifacts")[0].get("base64") | |
| elif "titan-image" in model_id: | |
| # For Titan model | |
| image_base64 = response_body.get("images")[0] | |
| elif "claude" in model_id: | |
| # For Claude model | |
| for content in response_body.get("content", []): | |
| if content.get("type") == "image": | |
| image_base64 = content.get("source", "") | |
| break | |
| else: | |
| # Try to find image data in response | |
| if "images" in response_body and len(response_body["images"]) > 0: | |
| image_base64 = response_body["images"][0] | |
| elif "image" in response_body: | |
| image_base64 = response_body["image"] | |
| else: | |
| print( | |
| f"Couldn't find image data in response: {json.dumps(response_body, indent=2)}" | |
| ) | |
| return None | |
| # Create output directory if it doesn't exist | |
| os.makedirs(output_dir, exist_ok=True) | |
| # Save the image | |
| image_data = base64.b64decode(image_base64) | |
| image = Image.open(io.BytesIO(image_data)) | |
| # Generate filename with timestamp | |
| timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") | |
| filename = f"{output_dir}/image_{timestamp}.png" | |
| image.save(filename) | |
| print(f"Image generated and saved to {filename}") | |
| return filename | |
| except Exception as e: | |
| print(f"Error generating image: {str(e)}") | |
| return None | |
| def list_available_models(): | |
| """List available Bedrock models""" | |
| try: | |
| bedrock = boto3.client("bedrock", region_name="us-east-1") | |
| response = bedrock.list_foundation_models() | |
| print("Available Bedrock models:") | |
| print("-" * 50) | |
| image_models = [] | |
| for model in response.get("modelSummaries", []): | |
| model_id = model.get("modelId") | |
| if model.get("inferenceTypesSupported") and "ON_DEMAND" in model.get( | |
| "inferenceTypesSupported" | |
| ): | |
| if "image" in model.get("outputModalities", []): | |
| image_models.append(model_id) | |
| print(f"✓ {model_id} (Image Generation)") | |
| print("-" * 50) | |
| print(f"Found {len(image_models)} models supporting image generation") | |
| return image_models | |
| except Exception as e: | |
| print(f"Error listing models: {str(e)}") | |
| return [] | |
| def main(): | |
| parser = argparse.ArgumentParser( | |
| description="Generate images from text using AWS Bedrock" | |
| ) | |
| parser.add_argument("prompt", nargs="?", help="Text prompt to generate image from") | |
| parser.add_argument( | |
| "--model", | |
| default="amazon.titan-image-generator-v1", | |
| help="Bedrock model ID (default: amazon.titan-image-generator-v1)", | |
| ) | |
| parser.add_argument( | |
| "--output-dir", | |
| default="generated_images", | |
| help="Directory to save generated images (default: generated_images)", | |
| ) | |
| parser.add_argument( | |
| "--list-models", | |
| action="store_true", | |
| help="List available Bedrock models for image generation", | |
| ) | |
| args = parser.parse_args() | |
| if args.list_models: | |
| list_available_models() | |
| elif args.prompt: | |
| generate_image(args.prompt, args.model, args.output_dir) | |
| else: | |
| parser.print_help() | |
| if __name__ == "__main__": | |
| main() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment