import torch
from diffusers import DiffusionPipeline
import os

# --- 1. 模型加载 ---
model_name = "/home/zwq/model/Qwen-Image"

print("正在加载模型... 这可能需要一些时间。")

pipe = DiffusionPipeline.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16, 
    device_map="balanced"           
)
print("模型加载成功！")


prompts_list = [
    '''A cute capybara wearing a top hat, sitting in a library.''',
    '''一幅中国水墨画，描绘了雨后云雾缭绕的黄山，几棵松树从岩石缝中伸出，意境悠远。''',
    '''A futuristic cityscape at night, with flying cars and neon signs that read "QWEN-IMAGE" and "通义千问", high-tech style.''',
    '''A coffee shop entrance features a chalkboard sign reading "Qwen Coffee 😊 $2 per cup," with a neon light beside it displaying "通义千问". Next to it hangs a poster showing a beautiful Chinese woman, and beneath the poster is written "π≈3.1415926-53589793-23846264-33832795-02384197". Ultra HD, 4K, cinematic composition'''
    '''A breathtaking tropical rainforest at dawn, with golden sunlight filtering through dense, vibrant green canopy, illuminating a crystal-clear waterfall cascading over moss-covered rocks. Exotic orchids and ferns in vivid reds and purples dot the scene, with delicate mist rising from the ground. Ultra-realistic style, hyper-detailed textures, cinematic lighting, 8K resolution.''',
    '''A hyper-detailed portrait of an elderly Chinese woman with intricate facial wrinkles and wise eyes, wearing an embroidered silk qipao with floral patterns. She stands in a traditional courtyard with blooming peonies and a wooden pavilion under soft, diffused moonlight. Realistic skin texture, subtle shadows, 8K resolution, cinematic depth of field.''',
    '''An abstract artwork inspired by Piet Mondrian, featuring bold geometric shapes—triangles, circles, and squares—in vibrant red, electric blue, and golden yellow, with smooth color gradients and subtle grainy textures. The background has a faint metallic sheen, creating a modern, dynamic composition, 4K ultra HD, sharp edges.''',
    '''A futuristic cyberpunk city at midnight, with towering skyscrapers clad in reflective glass, glowing with neon signs in pink, cyan, and purple, displaying "AI Revolution" and "通义千问". Hovering vehicles streak through the air, and holographic billboards cast dynamic reflections on wet streets after rain. Hyper-detailed, 8K resolution, cinematic atmosphere.''',
    '''A close-up of a snow leopard perched on a jagged, snow-dusted cliff at sunrise, its thick, silvery fur glistening with intricate patterns. The background features rugged Himalayan peaks under a fiery orange sky, with faint wisps of clouds. Ultra-realistic fur texture, dramatic lighting, 4K ultra HD.''',
    '''A colossal sci-fi mothership orbiting a ringed gas giant, its surface etched with intricate circuitry and glowing blue panels. Small fighter drones swarm around it, and a vibrant nebula in shades of purple and magenta fills the starry background. Hyper-detailed metallic textures, cinematic lighting, 8K resolution.''',
    '''A cozy still life on a rustic oak table, featuring a translucent glass teapot steaming with jasmine tea, a porcelain cup with delicate floral patterns, and a plate of freshly baked matcha cookies. Soft morning light streams through a lace-curtained window, casting intricate shadows. Photorealistic, 4K resolution, warm tones.''',
    '''A traditional Chinese ink painting of a misty bamboo forest in early spring, with slender bamboo stalks swaying gently in the breeze. A small stone bridge arches over a tranquil stream, and distant mountains fade into soft gray fog. Monochromatic with subtle green accents, ultra-detailed brushstrokes, 8K resolution.'''
]


positive_magic = "Ultra HD, 4K, cinematic composition."
negative_prompt = " "
width, height = (1664, 928)  
num_inference_steps = 50
true_cfg_scale = 4.0
base_seed = 42  

output_dir = "generated_images_GPU"
os.makedirs(output_dir, exist_ok=True)
print(f"所有图片将保存在 '{output_dir}' 文件夹中。")


for i, prompt in enumerate(prompts_list):
    
    print(f"\n正在处理第 {i+1}/{len(prompts_list)} 个prompt...")
    print(f"Prompt: '{prompt[:80]}...'") 


    generator = torch.Generator(device="cpu").manual_seed(base_seed + i) 
    
    image = pipe(
        prompt=prompt + " " + positive_magic, 
        negative_prompt=negative_prompt,
        width=width,
        height=height,
        num_inference_steps=num_inference_steps,
        true_cfg_scale=true_cfg_scale,
        generator=generator
    ).images[0]

    output_filename = f"prompt_{i}.png"
    output_path = os.path.join(output_dir, output_filename)
    
    image.save(output_path)
    print(f"图片已成功保存为: '{output_path}'")


print("\n所有prompt处理完毕！")