Flux Kontext API 快速开始

import requests
import time

class FluxKontextAPI:
    def __init__(self, api_key):
        self.api_key = api_key
        self.base_url = 'https://api.kie.ai/api/v1/flux/kontext'
        self.headers = {
            'Authorization': f'Bearer {api_key}',
            'Content-Type': 'application/json'
        }
    
    def generate_image(self, prompt, **options):
        data = {
            'prompt': prompt,
            'aspectRatio': options.get('aspectRatio', '16:9'),
            'model': options.get('model', 'flux-kontext-pro'),
            'enableTranslation': options.get('enableTranslation', True),
            'outputFormat': options.get('outputFormat', 'jpeg'),
            **options
        }
        
        response = requests.post(f'{self.base_url}/generate', 
                               headers=self.headers, json=data)
        result = response.json()
        
        if not response.ok or result.get('code') != 200:
            raise Exception(f"生成失败: {result.get('msg', '未知错误')}")
        
        return result['data']['taskId']
    
    def edit_image(self, prompt, input_image, **options):
        return self.generate_image(prompt, inputImage=input_image, **options)
    
    def wait_for_completion(self, task_id, max_wait_time=300):
        start_time = time.time()
        
        while time.time() - start_time < max_wait_time:
            status = self.get_task_status(task_id)
            success_flag = status['successFlag']
            
            if success_flag == 0:
                print("任务生成中，继续等待...")
            elif success_flag == 1:
                print("生成成功完成！")
                return status['response']
            elif success_flag == 2:
                create_error = status.get('errorMessage', '创建任务失败')
                print(f"创建任务失败: {create_error}")
                if status.get('errorCode'):
                    print(f"错误代码: {status['errorCode']}")
                raise Exception(create_error)
            elif success_flag == 3:
                generate_error = status.get('errorMessage', '生成失败')
                print(f"生成失败: {generate_error}")
                if status.get('errorCode'):
                    print(f"错误代码: {status['errorCode']}")
                raise Exception(generate_error)
            else:
                print(f"未知状态: {success_flag}")
                if status.get('errorMessage'):
                    print(f"错误信息: {status['errorMessage']}")
            
            time.sleep(3)  # 等待3秒
        
        raise Exception('生成超时')
    
    def get_task_status(self, task_id):
        response = requests.get(f'{self.base_url}/record-info?taskId={task_id}',
                              headers={'Authorization': f'Bearer {self.api_key}'})
        result = response.json()
        
        if not response.ok or result.get('code') != 200:
            raise Exception(f"查询状态失败: {result.get('msg', '未知错误')}")
        
        return result['data']

# 使用示例
def main():
    api = FluxKontextAPI('YOUR_API_KEY')
    
    try:
        # 文本转图像生成
        print('开始图像生成...')
        task_id = api.generate_image(
            '夜晚的未来主义城市景观，有霓虹灯和飞行汽车',
            aspectRatio='16:9',
            model='flux-kontext-max',
            promptUpsampling=True
        )
        
        # 等待完成
        print(f'任务ID: {task_id}。等待完成...')
        result = api.wait_for_completion(task_id)
        
        print('图像生成成功！')
        print(f'结果图像URL: {result["resultImageUrl"]}')
        print(f'原始图像URL（10分钟有效）: {result["originImageUrl"]}')
        
        # 图像编辑示例
        print('\n开始图像编辑...')
        edit_task_id = api.edit_image(
            '在天空中添加彩虹',
            result['resultImageUrl'],
            aspectRatio='16:9'
        )
        
        edit_result = api.wait_for_completion(edit_task_id)
        print('图像编辑成功！')
        print(f'编辑后图像URL: {edit_result["resultImageUrl"]}')
        
        # 下载图像
        download_image(result['resultImageUrl'], 'generated_image.jpg')
        download_image(edit_result['resultImageUrl'], 'edited_image.jpg')
        
    except Exception as error:
        print(f'错误: {error}')

def download_image(url, filename):
    response = requests.get(url)
    response.raise_for_status()
    
    with open(filename, 'wb') as f:
        f.write(response.content)
    print(f'已下载: {filename}')

if __name__ == '__main__':
    main()