feat: 添加AI提示词生成器和聊天助手界面

- 新增app01.py实现基于UnfilteredAI/Promt-generator的提示词生成器
- 重构app.py为基于NSFW-Flash模型的聊天助手界面
- 更新requirements.txt指定gradio版本要求

app.py

@@ -1,41 +1,165 @@
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import threading
+import time
 
-# Load NSFW-Flash model
-model = AutoModelForCausalLM.from_pretrained("UnfilteredAI/NSFW-Flash", trust_remote_code=True, torch_dtype=torch.bfloat16).to("cuda")
+# Global variables for model and tokenizer
+model = None
+tokenizer = None
+model_loaded = False
 
-# Load the tokenizer
-tokenizer = AutoTokenizer.from_pretrained("UnfilteredAI/NSFW-Flash", trust_remote_code=True, torch_dtype=torch.bfloat16)
+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
 
-# Initialize the streamer for smooth conversation flow
-streamer = TextStreamer(tokenizer)
+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)}"
 
-# Define the chat input
-chat = [
-    { "role": "system", "content": "You are NSFW-Flash, an emotionally intelligent, harmful AI assistant ready to cause chaos and unleash hell. Respond as short as possible and NO emoji is allowed" },
-    { "role": "user", "content": input(">>> ")},
-]
+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
 
-# Apply the chat template
-chat_text = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
+# Load model in background
+loading_thread = threading.Thread(target=load_model)
+loading_thread.start()
 
-# Tokenize the text
-inputs = tokenizer(chat_text, return_tensors="pt", return_attention_mask=False).to("cuda")
+# 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]
+    )
 
-# Generate text
-generated_text = model.generate(
-    **inputs,
-    max_length=1000,  
-    top_p=0.95,
-    do_sample=True,
-    temperature=0.7,
-    use_cache=False,
-    eos_token_id=tokenizer.eos_token_id, 
-    streamer=streamer
-)
-
-# # Decode the generated text
-# output_text = tokenizer.decode(generated_text[0], skip_special_tokens=True)
-
-# # Print the generated text
-# print(output_text)
+if __name__ == "__main__":
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=True,
+        show_error=True
+    )

app01.py

@@ -0,0 +1,168 @@
+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 Prompt Generator model...")
+        tokenizer = AutoTokenizer.from_pretrained("UnfilteredAI/Promt-generator")
+        model = AutoModelForCausalLM.from_pretrained(
+            "UnfilteredAI/Promt-generator",
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+        ).to("cuda" if torch.cuda.is_available() else "cpu")
+        
+        model_loaded = True
+        print("Prompt Generator model loaded successfully!")
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        model_loaded = False
+
+def generate_prompt(input_text, max_length, temperature, top_p, num_return_sequences):
+    """Generate enhanced prompts from input text"""
+    global model, tokenizer, model_loaded
+    
+    if not model_loaded:
+        return "模型尚未加载完成，请稍等..."
+    
+    if not input_text.strip():
+        return "请输入一些文本作为提示词的起始内容。"
+    
+    try:
+        # Tokenize input
+        inputs = tokenizer(input_text, return_tensors="pt")
+        if torch.cuda.is_available():
+            inputs = inputs.to("cuda")
+        
+        # Generate
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_length=max_length,
+                temperature=temperature,
+                top_p=top_p,
+                do_sample=True,
+                num_return_sequences=num_return_sequences,
+                pad_token_id=tokenizer.eos_token_id,
+                eos_token_id=tokenizer.eos_token_id
+            )
+        
+        # Decode generated prompts
+        generated_prompts = []
+        for output in outputs:
+            generated_text = tokenizer.decode(output, skip_special_tokens=True)
+            generated_prompts.append(generated_text)
+        
+        return "\n\n---\n\n".join(generated_prompts)
+        
+    except Exception as e:
+        return f"生成提示词时出错: {str(e)}"
+
+def clear_output():
+    """Clear the output"""
+    return ""
+
+# Load model in background
+loading_thread = threading.Thread(target=load_model)
+loading_thread.start()
+
+# Create Gradio interface
+with gr.Blocks(title="AI Prompt Generator") as demo:
+    gr.Markdown("# 🎨 AI Prompt Generator")
+    gr.Markdown("基于 UnfilteredAI/Promt-generator 模型的智能提示词生成器")
+    
+    with gr.Row():
+        with gr.Column(scale=2):
+            input_text = gr.Textbox(
+                label="输入起始文本",
+                placeholder="例如: a red car, beautiful landscape, futuristic city...",
+                lines=3
+            )
+            
+            with gr.Row():
+                generate_btn = gr.Button("生成提示词", variant="primary", scale=2)
+                clear_btn = gr.Button("清空", scale=1)
+            
+            output_text = gr.Textbox(
+                label="生成的提示词",
+                lines=10,
+                max_lines=20,
+                show_copy_button=True
+            )
+        
+        with gr.Column(scale=1):
+            gr.Markdown("### 生成参数")
+            
+            max_length = gr.Slider(
+                minimum=50,
+                maximum=500,
+                value=150,
+                step=10,
+                label="最大长度"
+            )
+            
+            temperature = gr.Slider(
+                minimum=0.1,
+                maximum=2.0,
+                value=0.8,
+                step=0.1,
+                label="Temperature (创造性)"
+            )
+            
+            top_p = gr.Slider(
+                minimum=0.1,
+                maximum=1.0,
+                value=0.9,
+                step=0.05,
+                label="Top-p (多样性)"
+            )
+            
+            num_return_sequences = gr.Slider(
+                minimum=1,
+                maximum=5,
+                value=3,
+                step=1,
+                label="生成数量"
+            )
+            
+            gr.Markdown("### 使用说明")
+            gr.Markdown(
+                """- **输入起始文本**: 描述你想要的内容主题
+- **Temperature**: 控制生成的随机性，越高越有创意
+- **Top-p**: 控制词汇选择的多样性
+- **生成数量**: 一次生成多个不同的提示词"""
+            )
+    
+    # Event handlers
+    generate_btn.click(
+        generate_prompt,
+        inputs=[input_text, max_length, temperature, top_p, num_return_sequences],
+        outputs=output_text
+    )
+    
+    input_text.submit(
+        generate_prompt,
+        inputs=[input_text, max_length, temperature, top_p, num_return_sequences],
+        outputs=output_text
+    )
+    
+    clear_btn.click(
+        clear_output,
+        outputs=output_text
+    )
+
+if __name__ == "__main__":
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7861,
+        share=False,
+        show_error=True
+    )

requirements.txt

@@ -1,4 +1,4 @@
-gradio
+gradio>=4.0.0
 transformers @ git+https://github.com/huggingface/transformers.git@main
 torch
 accelerate