app.py

import os
import requests
import gradio as gr

HF_MODEL = "TOPAI-Network/nsfw_chat_0124"
HF_TOKEN = os.getenv("HF_TOKEN")
API_URL = f"https://router.huggingface.co/hf-inference/models/{HF_MODEL}"
HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}

# 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."

    # 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}
    }

    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])

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