#### 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!")