Spaces:
Running
Running
| import pickle | |
| from pathlib import Path | |
| import os | |
| import requests | |
| import tqdm | |
| def generate_image(prompt, api_key: str, n=1, size="1024x1024"): | |
| """ | |
| Generates an image using OpenAI's DALL-E API. | |
| :param prompt: The prompt to generate an image for. | |
| :param api_key: Your OpenAI API key. | |
| :param n: The number of images to generate. | |
| :param size: The size of the generated images. | |
| :return: The response from the API call. | |
| """ | |
| headers = { | |
| "Content-Type": "application/json", | |
| "Authorization": f"Bearer {api_key}" | |
| } | |
| payload = { | |
| "model": "dall-e-3", | |
| "prompt": prompt, | |
| "n": n, | |
| "size": size | |
| } | |
| response = requests.post("https://api.openai.com/v1/images/generations", | |
| headers=headers, | |
| json=payload) | |
| return response | |
| def get_prompt(persona: dict): | |
| prompt = f"""A character named {persona['name']} is a {persona['age']} years old. The basic information of this character is as follows: | |
| - Name: {persona['name']} | |
| - Age: {persona['age']} | |
| - Gender: {persona['gender']} | |
| - Routine: {persona['routine']} | |
| - Personality: {', '.join(persona['personality'])} | |
| - Occupation: {persona['occupation']} | |
| - Thoughts: {persona['thoughts']} | |
| - Lifestyle: {persona['lifestyle']} | |
| Please generate a full-face photo of this character. | |
| """ | |
| return prompt | |
| def get_persona_avatar_bytes(persona: dict, api_key: str) -> bytes: | |
| prompt = get_prompt(persona) | |
| response = generate_image(prompt, api_key) | |
| if response.status_code == 200: | |
| # 这里可以添加代码来处理响应体，例如保存图像或进一步的处理 | |
| image_data = response.json()['data'] | |
| assert len(image_data) == 1 | |
| image_url = image_data[0]['url'] | |
| # download image | |
| image_response = requests.get(image_url) | |
| image = image_response.content | |
| return image | |
| else: | |
| print(f"Failed to generate image: {response.status_code} - {response.text}") | |
| def generate_for_cache(api_key: str): | |
| root_path = Path("static/exist_characters") | |
| character_dirs = os.listdir(root_path) | |
| character_dirs = [root_path / character_dir for character_dir in character_dirs] | |
| for character_dir in tqdm.tqdm(character_dirs): | |
| character_path = character_dir / f"{character_dir.name}.pkl" | |
| character_data = pickle.load(open(character_path, 'rb')) | |
| image_path = character_dir / f"avatar.jpg" | |
| image_bytes = get_persona_avatar_bytes(character_data['persona'], api_key) | |
| if image_bytes: | |
| with open(image_path, 'wb') as f: | |
| f.write(image_bytes) | |
| if __name__ == "__main__": | |
| api_key = os.getenv("OPENAI_API_KEY") | |
| generate_for_cache(api_key) | |