# app.py - 融合版：fuse_lora 加载滩头木版年画（零报错，2025稳）
import gradio as gr
from diffusers import StableDiffusionXLPipeline, StableDiffusionImg2ImgPipeline
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
import numpy as np

# 你的 LoRA 配置
LORA_MODEL_ID = "danyudan/chinese-tantou-wooden-nianhua"
LORA_FILENAME = "滩头木版年画.safetensors"

# 设备检测
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# 加载 base 模型：Segmind-Vega（SDXL 兼容，适合年画色彩）
pipe_txt2img = StableDiffusionXLPipeline.from_pretrained(
    "segmind/Segmind-Vega",
    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
    use_safetensors=True,
    variant="fp16" if device == "cuda" else None
).to(device)

pipe_img2img = StableDiffusionImg2ImgPipeline.from_pretrained(
    "segmind/Segmind-Vega",
    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
    use_safetensors=True,
    variant="fp16" if device == "cuda" else None
).to(device)

# 加载并融合 LoRA（用 fuse_lora，无需 adapter 名，兼容旧 LoRA）
lora_scale = 0.8  # 权重：0.6-1.0，0.8 最适合滩头风格
status_msg = "LoRA 加载中..."
try:
    lora_path = hf_hub_download(repo_id=LORA_MODEL_ID, filename=LORA_FILENAME)
    # 加载（prefix=None 消警告）
    pipe_txt2img.load_lora_weights(lora_path, weight_name=LORA_FILENAME, low_cpu=False)
    pipe_img2img.load_lora_weights(lora_path, weight_name=LORA_FILENAME, low_cpu=False)
    # 融合（直接加权，无 set_adapters）
    pipe_txt2img.fuse_lora(lora_scale=lora_scale)
    pipe_img2img.fuse_lora(lora_scale=lora_scale)
    status_msg = f"✅ 滩头木版年画 LoRA 已融合 (scale={lora_scale})"
    print(status_msg)
except Exception as e:
    status_msg = f"❌ LoRA 加载失败：{str(e)}。用基础模型（效果稍差）。"
    print(status_msg)

def generate_nianhua(prompt, negative_prompt, seed, steps=25):
    generator = torch.manual_seed(int(seed))
    full_prompt = prompt + ", tantou nianhua, tan tou woodblock new year print, beach-head style, 滩头木版年画, vibrant colors, bold black outlines, symmetrical, auspicious symbols, paper-cut texture, traditional chinese folk art"
    image = pipe_txt2img(
        prompt=full_prompt,
        negative_prompt=negative_prompt + ", blurry, low quality, modern, deformed",
        num_inference_steps=steps,
        guidance_scale=7.5,
        generator=generator,
        height=768, width=768,
        output_type="pil"
    ).images[0]
    return image

def transform_to_nianhua(face_image, strength=0.75, seed=42):
    if face_image is None:
        return None
    init_image = face_image.resize((768, 768))
    generator = torch.manual_seed(int(seed))
    full_prompt = "a person in tantou nianhua style, peeking from red door, traditional chinese costume, bold black outlines, vibrant colors, beach-head woodblock print, 滩头木版年画, symmetrical, auspicious pattern"
    result = pipe_img2img(
        prompt=full_prompt,
        negative_prompt="blurry, modern, realistic, deformed, extra limbs",
        image=init_image,
        strength=strength,
        num_inference_steps=30,
        guidance_scale=8,
        generator=generator,
        output_type="pil"
    ).images[0]
    return result

# Gradio 界面（状态用 Markdown 显示）
with gr.Blocks(title="AI滩头木版年画工坊") as demo:
    gr.Markdown("# 🎊 滩头木版年画生成 + 一键化身探头娃")
    status = gr.Markdown(value=status_msg)  # 动态状态
    gr.Markdown("专属探头年画！上传自拍变门神/胖娃。使用提示词如 '门神探头' 触发风格。")
    
    with gr.Tabs():
        with gr.Tab("📝 文字生成滩头年画"):
            prompt = gr.Textbox(label="正向提示词", value="门神探头守护家园, 喜庆红金, 滩头风格", lines=3)
            neg = gr.Textbox(label="反向提示词", value="模糊, 现代, 恐怖, 黑白", lines=2)
            seed1 = gr.Slider(0, 2**32-1, value=888, label="种子")
            steps = gr.Slider(20, 50, value=25, label="步数")
            img_out1 = gr.Image(label="生成结果")
            btn1 = gr.Button("生成年画", variant="primary")
            btn1.click(generate_nianhua, inputs=[prompt, neg, seed1, steps], outputs=img_out1)
            
        with gr.Tab("🤳 自拍变身探头人物"):
            gr.Markdown("### 正脸照 → 秒变滩头胖娃/门神（推荐清晰照）")
            face_input = gr.Image(label="上传自拍", type="pil")
            strength = gr.Slider(0.3, 1.0, value=0.75, label="年画化强度")
            seed2 = gr.Slider(0, 2**32-1, value=2025, label="种子")
            img_out2 = gr.Image(label="探头版你")
            btn2 = gr.Button("一键变身", variant="primary")
            btn2.click(transform_to_nianhua, inputs=[face_input, strength, seed2], outputs=img_out2)
    
    gr.Markdown("---\nMade with ❤️ | 2026 滩头年画，探头见喜！")

if __name__ == "__main__":
    demo.launch()