Spaces:
Running
Running
File size: 2,765 Bytes
8f3b56b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
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)
|