import base64
import time
import os
from gradio_client import Client
from PIL import Image
from io import BytesIO
from tqdm import tqdm  # 确保只使用标准输出版本的tqdm

# 创建 Gradio 客户端
client = Client("stabilityai/stable-diffusion-3-medium")

def generate_image(prompt, negative_prompt, save_path):
    try:
        # 调用 Gradio API 生成图像
        result = client.predict(
            prompt=prompt,
            negative_prompt=negative_prompt,
            seed=0,
            randomize_seed=True,
            width=1024,
            height=1024,
            guidance_scale=5,
            num_inference_steps=28,
            api_name="/infer"
        )

        # 打印 result 以查看其结构
        print(f"Result structure: {result}")

        # 确保 result 是 base64 编码的图像或图像 URL
        if isinstance(result, tuple):
            image_data = result[0]  # 假设图像数据在第一个元素中
        else:
            image_data = result

        # 如果是 base64 编码的数据
        if isinstance(image_data, str) and image_data.startswith('data:image'):
            image_data = image_data.split(",")[1]  # 移除前缀部分

        # 解码 base64 图像数据
        image = Image.open(BytesIO(base64.b64decode(image_data)))

        # 保存图像到本地指定目录
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        image.save(save_path, "PNG")
        print(f"Image saved as {save_path}")
    except Exception as e:
        print(f"Error generating image for prompt '{prompt}': {e}")

def main():
    input_file = r"D:\test\1.txt"  # 替换为你的TXT文件路径
    output_dir = r"D:\img"  # 替换为你想要保存图像的目录

    # 读取 TXT 文件中的每一行作为提示词
    with open(input_file, "r", encoding='utf-8') as file:
        prompts = file.readlines()

    negative_prompt = "default negative prompt"  # 可以替换为你想要的负提示词

    for i, prompt in enumerate(tqdm(prompts, ncols=80)):
        prompt = prompt.strip()  # 移除行末尾的换行符和空白符
        if not prompt:
            continue  # 跳过空行

        save_path = os.path.join(output_dir, f"Fake_Stable_Diffusion_3_medium_{i+1}.png")
        generate_image(prompt, negative_prompt, save_path)
        
        # 等待10秒
        time.sleep(10)

if __name__ == "__main__":
    main()