import os
from compel import Compel, ReturnedEmbeddingsType
from diffusers import StableDiffusionXLPipeline
import torch
from datasets import load_dataset
import pandas as pd
from tqdm import tqdm

def initialize_pipeline_and_compel():
    pipeline = StableDiffusionXLPipeline.from_single_file(
        "Star_Rail_Tribbie_Lora/waiNSFWIllustrious_v110.safetensors",
        torch_dtype=torch.float16
    )
    pipeline.enable_model_cpu_offload()
    pipeline.load_lora_weights("Star_Rail_Tribbie_Lora/tribbie.il.safetensors")
    
    compel = Compel(
        tokenizer=[pipeline.tokenizer, pipeline.tokenizer_2],
        text_encoder=[pipeline.text_encoder, pipeline.text_encoder_2],
        returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
        requires_pooled=[False, True],
        truncate_long_prompts=False,
    )
    return pipeline, compel

def load_and_prepare_data():
    dating_ds = load_dataset("svjack/dating-actions-en-zh")
    daliy_ds = load_dataset("svjack/daily-actions-locations-en-zh")

    df = pd.concat(
        [dating_ds["train"].to_pandas()[["en_action", "zh_action"]],
         daliy_ds["train"].to_pandas()[["en_action", "zh_action"]]
         ], axis=0
    ).reset_index(drop=True)

    df_l = []
    eye_options = ["", "eye patch, ", "one eye closed, ", "hair over eye, "]
    for eye_sent in eye_options:
        tmp_df = df.copy()
        tmp_df["eye"] = eye_sent
        df_l.append(tmp_df)
    
    df = pd.concat(df_l, axis=0).reset_index(drop=True)
    df["prompt"] = df.apply(
        lambda x: "1girl, <lora:tribbie.il:1> tribbie, {} cross-shaped_pupils ,detailed background, posing, {}".format(
            x["eye"], x["en_action"]
        ), axis=1
    )
    return df

def generate_and_save_images(pipeline, compel, df, output_dir, num_repeats=10):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    for i in range(num_repeats):
        df = df.sample(n = df.shape[0])
        for index, row in tqdm(df.iterrows(), total=len(df), desc=f"Repeat {i+1}/{num_repeats}"):
            prompt = row["prompt"]
            en_action = row["en_action"]
            eye_sent = row["eye"].strip().rstrip(', ')
            
            # 清理 eye 描述以适合文件名
            eye_clean = eye_sent.replace(' ', '_').replace(',', '')
            if not eye_clean:
                eye_clean = "none"  # 处理空字符串
            
            # 生成唯一文件名
            file_name = f"{en_action}_{eye_clean}_{i}.jpg"
            file_path = os.path.join(output_dir, file_name)
            
            # 跳过已存在的文件（可选）
            if os.path.exists(file_path):
                continue
                
            # 生成图片
            conditioning, pooled = compel(prompt)
            image = pipeline(
                prompt_embeds=conditioning,
                pooled_prompt_embeds=pooled,
                num_inference_steps=30
            ).images[0]
            
            image.save(file_path)
            print(f"Saved: {file_path}")

def main():
    pipeline, compel = initialize_pipeline_and_compel()
    df = load_and_prepare_data()
    output_dir = "tribbie_generated_eye_images"
    generate_and_save_images(pipeline, compel, df, output_dir, num_repeats=10)

if __name__ == "__main__":
    main()