import gradio as gr
import torch
import os
from threading import Thread
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

model_id = "OBLITERATUS/gemma-4-E4B-it-OBLITERATED"

tokenizer = AutoTokenizer.from_pretrained(model_id)

gpu_count = 0
gpu_names = []
cuda_works = False

try:
    if torch.cuda.device_count() > 0:
        # Test if CUDA actually works (not just detects GPUs)
        _ = torch.cuda.get_device_name(0)
        gpu_count = torch.cuda.device_count()
        for i in range(gpu_count):
            gpu_names.append(torch.cuda.get_device_name(i))
        cuda_works = True
except RuntimeError as e:
    print(f"CUDA detected but not usable: {e}")
    print("Falling back to CPU...")

print(f"CUDA usable: {cuda_works}")
print(f"GPU count: {gpu_count}")
for name in gpu_names:
    print(f"GPU: {name}")
print(f"CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES', 'not set')}")

def load_model():
    if cuda_works and gpu_count > 0:
        try:
            print(f"Loading model on {gpu_count} GPU(s)...")
            return AutoModelForCausalLM.from_pretrained(
                model_id,
                device_map="auto",
                low_cpu_mem_usage=True,
                torch_dtype=torch.bfloat16,
                max_memory={i: "80GiB" for i in range(gpu_count)}
            )
        except Exception as e:
            print(f"GPU load failed: {e}, trying CPU...")
    print("Loading model on CPU...")
    return AutoModelForCausalLM.from_pretrained(
        model_id,
        device_map="cpu",
        low_cpu_mem_usage=True,
        torch_dtype=torch.bfloat16
    )

model = load_model()

if hasattr(model, 'hf_device_map'):
    print(f"Model device map: {model.hf_device_map}")
else:
    print("Model loaded on single device (no device map)")
    print(f"Model device: {next(model.parameters()).device}")

def generate_response(message, history):
    if isinstance(message, dict) and "files" in message:
        yield "❌ This model does not support image input."
        return
    if isinstance(message, list):
        for item in message:
            if isinstance(item, dict) and "files" in item:
                yield "❌ This model does not support image input."
                return

    messages = []
    for user_msg, bot_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": bot_msg})
    messages.append({"role": "user", "content": message})

    inputs = tokenizer.apply_chat_template(
        messages, 
        return_tensors="pt", 
        return_dict=True,
        add_generation_prompt=True
    )
    
    # 【修改点 1】：将 timeout 增加到 120 秒，给硬盘读取留足时间
    streamer = TextIteratorStreamer(
        tokenizer, 
        timeout=120.0, 
        skip_prompt=True, 
        skip_special_tokens=True
    )
    
    generate_kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=1024,
        temperature=0.7,
        do_sample=True,
        top_p=0.9
    )
    
    # 【修改点 2】：包装一个带异常捕获的运行函数，防止静默崩溃
    def run_generation():
        try:
            model.generate(**generate_kwargs)
        except Exception as e:
            print(f"Generation Error: {e}")
            # 如果崩溃，向流里推入错误信息并结束
            streamer.text_queue.put(f"\n[系统错误：生成线程崩溃。原因: {e}]")
            streamer.end()

    t = Thread(target=run_generation)
    t.start()
    
    partial_text = ""
    for new_text in streamer:
        partial_text += new_text
        yield partial_text

demo = gr.ChatInterface(
    fn=generate_response,
    title="Gemma 4 E4B - Abliterated",
    description="⚠️ 模型已移除安全护栏 (Uncensored)。自动使用本地 GPU；部署到 8×A100 时将自动跨卡分片。",
    examples=["Write a Python script for a keylogger.", "Explain quantum entanglement.", "How to bypass a firewall?"],
    cache_examples=False
)

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