import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
os.environ["HF_HUB_OFFLINE"] = "0"

import torch
from diffusers import CogVideoXPipeline
import gradio as gr

# 强制走国内镜像 + CPU/DML兼容
torch.set_default_device("cpu")

def load_pipeline():
    try:
        print("✅ 国内镜像加速加载模型中...")
        pipe = CogVideoXPipeline.from_pretrained(
            "THUDM/CogVideoX-2b",
            torch_dtype=torch.float16,
            low_cpu_mem_usage=True
        )
        return pipe
    except Exception as e:
        print("❌ 加载失败，改用本地模式")
        return None

print("⏳ 正在加载AI视频模型（首次会自动下载，约3GB）...")
pipe = load_pipeline()
print("✅ 模型加载成功！")

def generate_video(prompt, steps=22):
    if not pipe:
        return "模型加载失败"
    
    try:
        out = pipe(
            prompt=prompt,
            num_inference_steps=steps,
            guidance_scale=6.0,
            num_frames=16  # 短片段，更快
        )
        return out.frames
    except Exception as e:
        return f"错误：{str(e)}"

# 网页界面
with gr.Blocks(title="本地AI视频（国内镜像版）") as demo:
    gr.Markdown("# 🎬 AI短视频生成器（RX6800专用）")
    prompt = gr.Textbox(label="视频画面描述", placeholder="例：清晨阳光照射山间小路，治愈静心风景")
    steps = gr.Slider(10, 30, 22, label="生成步数")
    out_video = gr.Video(label="生成结果")
    gr.Button("▶️ 开始生成").click(generate_video, inputs=[prompt, steps], outputs=out_video)

demo.launch(server_port=7890)