import argparse
import os
import torch
from diffusers import DiffusionPipeline

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--prompt", type=str, required=True, help="输入提示词")
    parser.add_argument("--seed", type=int, default=42, help="随机种子 (0 表示随机)")
    parser.add_argument("--width", type=int, default=1664, help="图像宽度")
    parser.add_argument("--height", type=int, default=928, help="图像高度")
    args = parser.parse_args()

    model_name = "Qwen/Qwen-Image"

    # Load the pipeline
    if torch.cuda.is_available():
        torch_dtype = torch.bfloat16
        device = "cuda"
    else:
        torch_dtype = torch.float32
        device = "cpu"

    pipe = DiffusionPipeline.from_pretrained(model_name, torch_dtype=torch_dtype)
    pipe = pipe.to(device)

    positive_magic = {
        "en": ", Ultra HD, 4K, cinematic composition.",  # for english prompt
        "zh": ", 超清，4K，电影级构图."  # for chinese prompt
    }

    # 在这里注入了固定的负面提示词
    negative_prompt = "political figures, Chinese political figures, porn, nsfw"

    # 生成器（0 = 随机种子，不固定）
    generator = None if args.seed == 0 else torch.Generator(device=device).manual_seed(args.seed)

    # Generate image
    image = pipe(
        prompt=args.prompt + positive_magic["en"],
        negative_prompt=negative_prompt,
        width=args.width,
        height=args.height,
        num_inference_steps=50,
        true_cfg_scale=4.0,
        generator=generator
    ).images[0]

    # 输出目录
    output_dir = os.path.abspath("output")
    os.makedirs(output_dir, exist_ok=True)

    # 固定文件名
    output_path = os.path.join(output_dir, "output.png")
    image.save(output_path)

    seed_info = "random" if args.seed == 0 else args.seed
    print(f"✅ Image saved at: {output_path} (seed={seed_info}, size={args.width}x{args.height})")

if __name__ == "__main__":
    main()