DynamicScene/DynamicGeneration

1import pickle2from pathlib import Path3import os4import requests5import tqdm6 7def generate_image(prompt, api_key: str, n=1, size="1024x1024"):8    """9    Generates an image using OpenAI's DALL-E API.10 11    :param prompt: The prompt to generate an image for.12    :param api_key: Your OpenAI API key.13    :param n: The number of images to generate.14    :param size: The size of the generated images.15    :return: The response from the API call.16    """17    headers = {18        "Content-Type": "application/json",19        "Authorization": f"Bearer {api_key}"20    }21    payload = {22        "model": "dall-e-3",23        "prompt": prompt,24        "n": n,25        "size": size26    }27    28    response = requests.post("https://api.openai.com/v1/images/generations", 29                             headers=headers, 30                             json=payload)31    return response32 33def get_prompt(persona: dict):34    prompt = f"""A character named {persona['name']} is a {persona['age']} years old. The basic information of this character is as follows:35- Name: {persona['name']}36- Age: {persona['age']}37- Gender: {persona['gender']}38- Routine: {persona['routine']}39- Personality: {', '.join(persona['personality'])}40- Occupation: {persona['occupation']}41- Thoughts: {persona['thoughts']}42- Lifestyle: {persona['lifestyle']}43Please generate a full-face photo of this character.44"""45    return prompt46 47def get_persona_avatar_bytes(persona: dict, api_key: str) -> bytes:48    prompt = get_prompt(persona)49    response = generate_image(prompt, api_key)50    if response.status_code == 200:51        # 这里可以添加代码来处理响应体，例如保存图像或进一步的处理52        image_data = response.json()['data']53        assert len(image_data) == 154        image_url = image_data[0]['url']55        # download image56        image_response = requests.get(image_url)57        image = image_response.content58        return image59    else:60        print(f"Failed to generate image: {response.status_code} - {response.text}")61 62def generate_for_cache(api_key: str):63    root_path = Path("static/exist_characters")64    character_dirs = os.listdir(root_path)65    character_dirs = [root_path / character_dir for character_dir in character_dirs]66 67    for character_dir in tqdm.tqdm(character_dirs):68        character_path = character_dir / f"{character_dir.name}.pkl"69        character_data = pickle.load(open(character_path, 'rb'))70        image_path = character_dir / f"avatar.jpg"71        image_bytes = get_persona_avatar_bytes(character_data['persona'], api_key)72        if image_bytes:73            with open(image_path, 'wb') as f:74                f.write(image_bytes)75 76 77if __name__ == "__main__":78    api_key = os.getenv("OPENAI_API_KEY")79    generate_for_cache(api_key)80