svjack's picture
Update run_abd.py
e6037ba verified
#### huggingface-cli download svjack/Genshin-Impact-Portrait-with-Tags-Filtered-IID-Gender-SP --include "genshin_impact_ALBEDO_images_and_texts/*" --local-dir .
#### mv genshin_impact_ALBEDO_oil_painting_and_texts genshin_impact_ALBEDO_hidiffusion_images_and_texts
import os
from tqdm import tqdm
from gradio_client import Client, handle_file
from PIL import Image
from shutil import copy2
# 初始化客户端
client = Client("http://localhost:7860")
# 输入和输出文件夹路径
input_dir = "genshin_impact_ALBEDO_images_and_texts"
output_dir = "genshin_impact_ALBEDO_oil_painting_and_texts"
# 创建输出文件夹(如果不存在)
os.makedirs(output_dir, exist_ok=True)
# 获取所有.png文件
png_files = [f for f in os.listdir(input_dir) if f.endswith('.png')]
# 遍历每个.png文件
for png_file in tqdm(png_files, desc="Processing images"):
# 获取对应的.txt文件名
base_name = os.path.splitext(png_file)[0]
txt_file = base_name + ".txt"
# 检查.txt文件是否存在
txt_path = os.path.join(input_dir, txt_file)
if not os.path.exists(txt_path):
print(f"Warning: No corresponding .txt file found for {png_file}")
continue
# 读取.txt文件内容并修改prompt
with open(txt_path, 'r', encoding='utf-8') as f:
prompt = f.read()
####.replace(
#### "In this digital anime-style drawing,",
#### "In this realistic personal drawing,"
####)
prompt = ",".join(["In this realistic personal drawing"] + prompt.split(",")[1:])
# 调用API
result = client.predict(
input_image=handle_file(os.path.join(input_dir, png_file)),
prompt=prompt,
negative_prompt="blurry, ugly, duplicate, poorly drawn, deformed, mosaic",
seed=2718545171199645000,
guidance_scale=8.5,
scale=2,
controlnet_conditioning_scale=0.5,
strength=1,
controlnet_start=0,
controlnet_end=1,
guassian_sigma=2,
intensity_threshold=3,
api_name="/predict"
)
# 获取生成的图像路径
generated_image_path = result[-1][1]
# 使用PIL打开图像并保存到输出文件夹
output_image_path = os.path.join(output_dir, png_file)
with Image.open(generated_image_path) as img:
img.save(output_image_path)
# 保存修改后的prompt到.txt文件
output_txt_path = os.path.join(output_dir, txt_file)
with open(output_txt_path, 'w', encoding='utf-8') as f:
f.write(prompt)
print("Processing completed!")