Spaces:
Sleeping
Sleeping
import fal_client | |
import pandas as pd | |
from prompt_gen import prompt_gen | |
import requests | |
import os | |
from openai import OpenAI | |
import random | |
import shutil | |
from pathlib import Path | |
from PIL import Image | |
nv_prompt_file = pd.read_excel('汉服-女词库.xlsx') | |
na_prompt_file = pd.read_excel('汉服-男词库.xlsx') | |
nv_prompt = nv_prompt_file.to_string(index=False) | |
na_prompt = na_prompt_file.to_string(index=False) | |
save_directory = "downloads" | |
def prompt_nan(prompt): | |
client = OpenAI() | |
completion = client.chat.completions.create( | |
model="gpt-4o", | |
messages=[ | |
{"role": "system", | |
"content": "You are a helpful assistant.", }, | |
{"role": "user", | |
"content": "Please showcase the overall appearance of this Hanfu robe against a contrasting white " | |
"background. Highlight its intricate details and unique design elements, including" + prompt, | |
} | |
] | |
) | |
print("change prompt: ") | |
print(completion.choices[0].message.content) | |
return completion.choices[0].message.content | |
def pro_gen(advice, gender, index): | |
prompt = prompt_gen(advice, gender) | |
start_index = prompt.find("Begin") | |
if start_index == -1: | |
start_index = prompt.find("begin") | |
prompt__gen = "" | |
if start_index != -1: | |
start_index += len("Begin\n") | |
end_index = prompt.find("End") | |
if end_index != -1: | |
prompt__gen = prompt[start_index:end_index] | |
# if gender == "男": | |
# prompt__gen = prompt_nan(prompt__gen) | |
filename = os.path.join(save_directory, f"prompt_{index}.txt") | |
with open(filename, "w") as file: | |
file.write(prompt__gen) | |
# print(prompt__gen) | |
else: | |
print("No 'promptEnd' found after 'prompt'.") | |
else: | |
print("No 'prompt' found in the text.") | |
return prompt__gen | |
def generate(lora_path, prompt__gen, index): | |
# print(prompt__gen) | |
handler = fal_client.submit( | |
"fal-ai/fast-sdxl", | |
arguments={ | |
"prompt": prompt__gen, | |
"negative_prompt": "human, people, person, man, woman, child, model, face, head, eyes, hands, arms, legs, " | |
"feet, hair, portrait, worst quality, low quality, normal quality, lowres, signature, " | |
"watermark, jpeg artifacts, logo, monochrome, grayscale, ugly", | |
"image_size": "portrait_4_3", | |
"num_inference_steps": 28, | |
"guidance_scale": 7.5, | |
"num_images": 1, | |
"loras": [{"path": lora_path, "scale": 0.7}], | |
"embeddings": [], | |
"safety_checker_version": "v1", | |
"format": "jpeg" | |
}, | |
) | |
request_id = handler.request_id | |
result = fal_client.result("fal-ai/fast-sdxl", request_id) | |
image_index = index * 2 - 1 | |
for image in result['images']: | |
response = requests.get(image['url']) | |
if response.status_code == 200: | |
filename = os.path.join(save_directory, f"cloth_{image_index}.jpeg") | |
with open(filename, 'wb') as f: | |
f.write(response.content) | |
image_index += 1 | |
else: | |
print(f"Failed to download image from {image['url']}") | |
client = OpenAI() | |
completion = client.chat.completions.create( | |
model="gpt-4o", | |
messages=[ | |
{"role": "system", | |
"content": "You are a helpful assistant.", }, | |
{"role": "user", | |
"content": prompt__gen + "以上是一段对于一套汉服的描述,请根据描述内容对该套汉服进行介绍。要求以介绍的口吻输出内容", | |
} | |
] | |
) | |
cloth_intro = ("汉服,是汉民族的传统服饰。又称衣冠、衣裳、汉装。汉服是中国“衣冠上国”“礼仪之邦”“锦绣中华”的体现,承载了中国的染织绣等杰出" | |
"工艺和美学,传承了30多项中国非物质文化遗产以及受保护的中国工艺美术。\n") + completion.choices[ | |
0].message.content | |
filename = os.path.join(save_directory, f"cloth_intro_{index * 2 - 1}.txt") | |
with open(filename, "w") as file: | |
file.write(cloth_intro) | |
def convert_image_to_jpeg(input_path, output_path): | |
try: | |
image = Image.open(input_path) | |
if image.mode in ('RGBA', 'LA'): | |
image = image.convert('RGB') | |
image.save(output_path, 'JPEG') | |
except Exception as e: | |
print(f"转换图像时出错: {e}") | |
def pic_match(prompt__gen, cates, folder_path, intro_path, index): | |
client = OpenAI() | |
completion = client.chat.completions.create( | |
model="gpt-4o", | |
messages=[ | |
{"role": "system", | |
"content": "You are a helpful assistant.", }, | |
{"role": "user", | |
"content": prompt__gen + "以上是关于一套汉族服饰的描述,请根据描述内容从以下几种颜色中选择最符合描述的一种,可选颜色包括:" + cates | |
+ ". 仅需输出一种颜色名称,不要带任何符号", | |
} | |
] | |
) | |
print(f"Selected color: {completion.choices[0].message.content}") | |
folder_path = os.path.join(folder_path, completion.choices[0].message.content) | |
files = os.listdir(folder_path) | |
random_file = random.choice(files) | |
source_file_path = os.path.join(folder_path, random_file) | |
file_prefix, file_ext = os.path.splitext(random_file) | |
target_file_path = os.path.join(save_directory, f"cloth_{index * 2}.jpeg") | |
convert_image_to_jpeg(source_file_path, target_file_path) | |
file_extension = ".txt" | |
search_path = Path(intro_path) | |
for file in search_path.glob(f"{file_prefix}*{file_extension}"): | |
if file.is_file(): | |
with open(file, "r") as f: | |
content = f.read() | |
client = OpenAI() | |
completion = client.chat.completions.create( | |
model="gpt-4o", | |
messages=[ | |
{"role": "system", | |
"content": "You are a helpful assistant.", }, | |
{"role": "user", | |
"content": content + "以上是一段对于一套汉服的描述,请根据描述内容对该套汉服进行介绍。要求以介绍的口吻输出内容", | |
} | |
] | |
) | |
cloth_intro = ("汉服,是汉民族的传统服饰。又称衣冠、衣裳、汉装。汉服是中国“衣冠上国”“礼仪之邦”“锦绣中华”的体现,承载了中国的染织绣等杰出" | |
"工艺和美学,传承了30多项中国非物质文化遗产以及受保护的中国工艺美术。\n") + \ | |
completion.choices[0].message.content | |
filename = os.path.join(save_directory, f"cloth_intro_{index * 2}.txt") | |
with open(filename, "w") as file: | |
file.write(cloth_intro) | |
return target_file_path | |
def cloth_gen(gender): | |
cates = "Black, Blue, Green, Orange, Pink, Red, Violet, White, Yellow" | |
lora_path = "https://huggingface.co/PPSharks/PPSharksModels/resolve/main/NV.safetensors" | |
folder_path = "database/female" | |
intro_path = "database/female_intro" | |
if gender == "男": | |
lora_path = "https://huggingface.co/PPSharks/PPSharksModels/resolve/main/NA.safetensors" | |
cates = "Black, Blue, Green, Brown, Red, Violet" | |
folder_path = "database/male" | |
intro_path = "database/male_intro" | |
elif gender == "女": | |
lora_path = "https://huggingface.co/PPSharks/PPSharksModels/resolve/main/NV.safetensors" | |
cates = "Black, Blue, Green, Orange, Pink, Red, Violet, White, Yellow" | |
folder_path = "database/female" | |
intro_path = "database/female_intro" | |
cloth_image = [] | |
for i in range(1, 4): | |
with open(os.path.join(save_directory, f"prompt_{i}.txt"), "r") as file: | |
prompt__gen = file.read() | |
generate(lora_path, prompt__gen, i) | |
cloth_image.append(os.path.join(save_directory, f"cloth_{i*2-1}.jpeg")) | |
pic_path = pic_match(prompt__gen, cates, folder_path, intro_path, i) | |
cloth_image.append(pic_path) | |
with open(os.path.join(save_directory, f"cloth_intro_1.txt"), "r") as file: | |
cloth_intro = file.read() | |
return cloth_image, cloth_image[0], cloth_intro | |