Update app.py

app.py

@@ -1,48 +1,18 @@
-import os
-import requests
 import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-HF_MODEL = "TOPAI-Network/nsfw_chat_0124"
-HF_TOKEN = os.getenv("HF_TOKEN")
-API_URL = f"https://router.huggingface.co/hf-inference/models/TOPAI-Network/nsfw_chat_0124"
-HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}
+model_id = "mradermacher/nsfw_chat_0124-GGUF"
 
-# Safety / gate: simple age checkbox + confirmation of NSFW
-def generate(system_prompt, user_prompt, allow_nsfw):
-    if not allow_nsfw:
-        return "Bạn phải đánh dấu xác nhận để sử dụng model này (chứa nội dung nhạy cảm)."
-    if not HF_TOKEN:
-        return "HF_TOKEN chưa được cấu hình. Thêm token vào Secrets của Space."
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto", trust_remote_code=True)
 
-    # Kết hợp prompt (tùy chỉnh theo model)
-    payload = {
-        "inputs": f"{system_prompt}\nUser: {user_prompt}",
-        "options": {"use_cache": False, "wait_for_model": True}
-    }
+def chat(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    outputs = model.generate(**inputs, max_new_tokens=200)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    resp = requests.post(API_URL, headers=HEADERS, json=payload, timeout=120)
-    if resp.status_code != 200:
-        return f"Error {resp.status_code}: {resp.text}"
-    data = resp.json()
-    # Inference API trả về list hoặc dict tùy model; xử lý an toàn:
-    if isinstance(data, dict) and data.get("error"):
-        return "Model error: " + data["error"]
-    # Một số model trả về [{"generated_text": "..."}]
-    if isinstance(data, list) and "generated_text" in data[0]:
-        return data[0]["generated_text"]
-    # Nếu trả về chuỗi hoặc khác
-    return str(data)
-
-# UI
-with gr.Blocks(title="NSFW Chat (gated)") as demo:
-    gr.Markdown("## Cảnh báo: nội dung nhạy cảm\nModel này có thể tạo nội dung NSFW. Vui lòng xác nhận tuổi >=18 và chấp nhận rủi ro.")
-    system = gr.Textbox(value="You are a helpful assistant.", label="System prompt (tuỳ chọn)")
-    user_in = gr.Textbox(label="User prompt", placeholder="Nhập lời nhắc...")
-    allow = gr.Checkbox(label="Tôi xác nhận tôi 18+ và đồng ý xem nội dung nhạy cảm", value=False)
-    out = gr.Textbox(label="Model output")
-    btn = gr.Button("Gửi")
-
-    btn.click(fn=generate, inputs=[system, user_in, allow], outputs=[out])
+demo = gr.Interface(fn=chat, inputs="text", outputs="text", title="NSFW Chat 0124 (local)")
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()