app.00.py

import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import threading
import time

# Global variables for model and tokenizer
model = None
tokenizer = None
model_loaded = False

def load_model():
    """Load the model and tokenizer"""
    global model, tokenizer, model_loaded
    try:
        print("Loading model...")
        model = AutoModelForCausalLM.from_pretrained(
            "UnfilteredAI/NSFW-Flash", 
            trust_remote_code=True, 
            torch_dtype=torch.bfloat16
        ).to("cuda" if torch.cuda.is_available() else "cpu")
        
        tokenizer = AutoTokenizer.from_pretrained(
            "UnfilteredAI/NSFW-Flash", 
            trust_remote_code=True
        )
        
        model_loaded = True
        print("Model loaded successfully!")
    except Exception as e:
        print(f"Error loading model: {e}")
        model_loaded = False

def generate_response(message, history, temperature, max_length, top_p):
    """Generate response from the model"""
    global model, tokenizer, model_loaded
    
    if not model_loaded:
        return "模型尚未加载完成，请稍等..."
    
    try:
        # Build conversation history
        chat = [
            {"role": "system", "content": "You are NSFW-Flash, an AI assistant. Respond helpfully and appropriately."}
        ]
        
        # Add conversation history
        for user_msg, bot_msg in history:
            chat.append({"role": "user", "content": user_msg})
            if bot_msg:
                chat.append({"role": "assistant", "content": bot_msg})
        
        # Add current message
        chat.append({"role": "user", "content": message})
        
        # Apply chat template
        chat_text = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
        
        # Tokenize
        inputs = tokenizer(chat_text, return_tensors="pt", return_attention_mask=False)
        if torch.cuda.is_available():
            inputs = inputs.to("cuda")
        
        # Generate
        with torch.no_grad():
            generated = model.generate(
                **inputs,
                max_length=max_length,
                temperature=temperature,
                top_p=top_p,
                do_sample=True,
                use_cache=False,
                eos_token_id=tokenizer.eos_token_id,
                pad_token_id=tokenizer.eos_token_id
            )
        
        # Decode response
        response = tokenizer.decode(generated[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
        return response.strip()
        
    except Exception as e:
        return f"生成回复时出错: {str(e)}"

def chat_interface(message, history, temperature, max_length, top_p):
    """Chat interface for Gradio"""
    response = generate_response(message, history, temperature, max_length, top_p)
    history.append([message, response])
    return "", history

# Load model in background
loading_thread = threading.Thread(target=load_model)
loading_thread.start()

# Create Gradio interface
with gr.Blocks(title="AI Chat Assistant") as demo:
    gr.Markdown("# 🤖 AI Chat Assistant")
    gr.Markdown("基于 NSFW-Flash 模型的聊天助手")
    
    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(
                value=[],
                height=500,
                show_label=False
            )
            
            with gr.Row():
                msg = gr.Textbox(
                    placeholder="输入您的消息...",
                    show_label=False,
                    scale=4
                )
                send_btn = gr.Button("发送", scale=1)
            
            clear_btn = gr.Button("清空对话")
        
        with gr.Column(scale=1):
            gr.Markdown("### 参数设置")
            temperature = gr.Slider(
                minimum=0.1,
                maximum=2.0,
                value=0.7,
                step=0.1,
                label="Temperature"
            )
            max_length = gr.Slider(
                minimum=100,
                maximum=2000,
                value=1000,
                step=100,
                label="最大长度"
            )
            top_p = gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.95,
                step=0.05,
                label="Top-p"
            )
    
    # Event handlers
    send_btn.click(
        chat_interface,
        inputs=[msg, chatbot, temperature, max_length, top_p],
        outputs=[msg, chatbot]
    )
    
    msg.submit(
        chat_interface,
        inputs=[msg, chatbot, temperature, max_length, top_p],
        outputs=[msg, chatbot]
    )
    
    clear_btn.click(
        lambda: ([], ""),
        outputs=[chatbot, msg]
    )

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True,
        show_error=True
    )