app.py

"""
SuperGemma4-26B Uncensored GGUF - CPU Compatible
"""
import gradio as gr
from llama_cpp import Llama
import logging
from huggingface_hub import hf_hub_download
import os

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

MODEL_REPO = "Jiunsong/supergemma4-26b-uncensored-gguf-v2"
# Try Q2_K for smaller size and faster loading
MODEL_FILE = "supergemma4-26b-uncensored-Q2_K.gguf"

logger.info(f"Loading model: {MODEL_REPO}/{MODEL_FILE}")
logger.info(f"This may take 5-10 minutes for first load...")

llm = None

try:
    logger.info("Downloading model from HuggingFace...")
    model_path = hf_hub_download(
        repo_id=MODEL_REPO, 
        filename=MODEL_FILE, 
        repo_type="model",
        resume_download=True
    )
    logger.info(f"Model downloaded to: {model_path}")
    logger.info(f"Model file size: {os.path.getsize(model_path) / (1024**3):.2f} GB")
    
    logger.info("Loading model into memory...")
    llm = Llama(
        model_path=model_path,
        n_ctx=2048,  # Reduced context for faster inference
        n_threads=4,  # Reduced threads
        n_gpu_layers=0,  # CPU only
        verbose=True,
        n_batch=512
    )
    logger.info("✅ Model loaded successfully on CPU!")
    
except Exception as e:
    logger.error(f"❌ Error loading model: {str(e)}")
    logger.error(f"Full error: {repr(e)}")
    llm = None

def generate_text(prompt, max_tokens=500, temperature=0.7, top_p=0.9, top_k=40):
    if llm is None:
        return "❌ Error: Model not loaded. Check Space logs for details."
    try:
        logger.info(f"Generating: {prompt[:50]}...")
        response = llm(
            prompt, 
            max_tokens=int(max_tokens), 
            temperature=float(temperature), 
            top_p=float(top_p), 
            top_k=int(top_k), 
            stop=["</s>", "\n\n\n"], 
            echo=False
        )
        result = response['choices'][0]['text'].strip()
        logger.info(f"Generated {len(result)} characters")
        return result
    except Exception as e:
        logger.error(f"Generation error: {str(e)}")
        return f"Error: {str(e)}"

def generate_code(prompt, max_tokens=500, temperature=0.2, top_p=0.95):
    code_prompt = f"### Instruction:\nWrite code:\n{prompt}\n\n### Response:\n"
    return generate_text(code_prompt, max_tokens, temperature, top_p, 40)

def chat(message, history, max_tokens=500, temperature=0.7):
    if llm is None:
        return "❌ Error: Model not loaded"
    conversation = ""
    for user_msg, assistant_msg in history:
        conversation += f"User: {user_msg}\nAssistant: {assistant_msg}\n\n"
    conversation += f"User: {message}\nAssistant: "
    response = llm(
        conversation, 
        max_tokens=int(max_tokens), 
        temperature=float(temperature), 
        top_p=0.9, 
        top_k=40, 
        stop=["User:", "</s>"], 
        echo=False
    )
    return response['choices'][0]['text'].strip()

# Create UI
with gr.Blocks(title="SuperGemma4-26B Uncensored", theme=gr.themes.Soft()) as demo:
    gr.Markdown(f"""
    # 🚀 SuperGemma4-26B Uncensored (CPU)
    
    **Status**: {'✅ Model Loaded' if llm else '❌ Model Loading Failed'}
    
    26B parameter uncensored model running on CPU with Q2_K quantization
    
    ⚠️ **Note**: First load takes 5-10 minutes. Please be patient!
    """)
    
    if llm is None:
        gr.Markdown("""
        ### ⚠️ Model Loading Error
        
        The model failed to load. Possible reasons:
        1. Model file is still downloading (check Space logs)
        2. Insufficient memory
        3. Model file not found
        
        **Check the Logs tab** in your Space for detailed error messages.
        """)
    
    with gr.Tabs():
        with gr.Tab("💬 Chat"):
            chatbot = gr.Chatbot(height=400)
            msg = gr.Textbox(label="Message", placeholder="Ask anything...")
            with gr.Row():
                chat_max_tokens = gr.Slider(100, 1000, value=300, label="Max Tokens")
                chat_temperature = gr.Slider(0.1, 2.0, value=0.7, label="Temperature")
            with gr.Row():
                submit = gr.Button("Send", variant="primary")
                clear = gr.Button("Clear")
            
            def respond(message, chat_history, max_tokens, temperature):
                bot_message = chat(message, chat_history, max_tokens, temperature)
                chat_history.append((message, bot_message))
                return "", chat_history
            
            submit.click(respond, [msg, chatbot, chat_max_tokens, chat_temperature], [msg, chatbot])
            clear.click(lambda: None, None, chatbot, queue=False)
        
        with gr.Tab("💻 Generate Code"):
            with gr.Row():
                with gr.Column():
                    gen_prompt = gr.Textbox(label="Prompt", lines=5, placeholder="Write a Python function to...")
                    gen_max_tokens = gr.Slider(100, 1000, value=400, label="Max Tokens")
                    gen_temperature = gr.Slider(0.1, 1.0, value=0.2, label="Temperature")
                    gen_top_p = gr.Slider(0.1, 1.0, value=0.95, label="Top P")
                    gen_button = gr.Button("Generate", variant="primary")
                with gr.Column():
                    gen_output = gr.Textbox(label="Generated Code", lines=20)
            gen_button.click(generate_code, [gen_prompt, gen_max_tokens, gen_temperature, gen_top_p], gen_output)
        
        with gr.Tab("📝 Generate Text"):
            with gr.Row():
                with gr.Column():
                    text_prompt = gr.Textbox(label="Prompt", lines=5, placeholder="Write about...")
                    text_max_tokens = gr.Slider(100, 1000, value=400, label="Max Tokens")
                    text_temperature = gr.Slider(0.1, 2.0, value=0.7, label="Temperature")
                    text_top_p = gr.Slider(0.1, 1.0, value=0.9, label="Top P")
                    text_top_k = gr.Slider(1, 100, value=40, label="Top K")
                    text_button = gr.Button("Generate", variant="primary")
                with gr.Column():
                    text_output = gr.Textbox(label="Generated Text", lines=20)
            text_button.click(generate_text, [text_prompt, text_max_tokens, text_temperature, text_top_p, text_top_k], text_output)
    
    gr.Markdown("""
    ---
    **Model**: SuperGemma4-26B-Uncensored (Q2_K) | **Hardware**: CPU | **Powered by**: llama.cpp
    
    ⚠️ CPU inference is slower (~1-3 tokens/second). Be patient with responses.
    """)

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