app.py

"""Gemma-4-E4B-Turbo — ZeroGPU / llama.cpp / 32K context + multimodal"""
import spaces
import gradio as gr
from huggingface_hub import hf_hub_download
from pathlib import Path
import logging, sys, os, json, tempfile, time

logging.basicConfig(level=logging.INFO, stream=sys.stdout)
logger = logging.getLogger(__name__)

MODEL_REPO = "HauhauCS/Gemma-4-E4B-Uncensored-HauhauCS-Aggressive"
MODEL_FILE = "Gemma-4-E4B-Uncensored-HauhauCS-Aggressive-Q5_K_P.gguf"
MMPROJ_FILE = "mmproj-Gemma-4-E4B-Uncensored-HauhauCS-Aggressive-Q5_K_P.gguf"

class ModelManager:
    """Lazy-loading model singleton with error containment."""
    def __init__(self):
        self._llm = None
        self._has_mmproj = False
        self._mmproj_path = None
        self._model_path = None
        self._ready = False

    def _download(self):
        """Download model files once."""
        if self._model_path:
            return
        logger.info(f"Downloading {MODEL_FILE} from {MODEL_REPO}...")
        self._model_path = hf_hub_download(
            repo_id=MODEL_REPO, filename=MODEL_FILE, resume_download=True
        )
        try:
            self._mmproj_path = hf_hub_download(
                repo_id=MODEL_REPO, filename=MMPROJ_FILE, resume_download=True
            )
            self._has_mmproj = True
            logger.info("mmproj found — multimodal enabled")
        except Exception:
            self._has_mmproj = False
            logger.info("No mmproj — text-only mode")

    @spaces.GPU(duration=300)
    def load(self):
        """Load llama model on first call (GPU-backed)."""
        if self._ready:
            return self._llm
        self._download()
        logger.info("Loading model into GPU...")
        from llama_cpp import Llama
        kwargs = {
            "model_path": self._model_path,
            "n_gpu_layers": -1,  # Offload ALL layers to GPU
            "n_ctx": 32768,       # 32K context
            "n_threads": 8,
            "verbose": False,
            "use_mmap": True,
        }
        if self._has_mmproj:
            kwargs["mmproj"] = self._mmproj_path
        self._llm = Llama(**kwargs)
        self._ready = True
        logger.info("Model loaded and ready")
        return self._llm

model = ModelManager()

@spaces.GPU(duration=300)
def generate(prompt, max_tokens=1024, temperature=0.7, top_p=0.9, repeat_penalty=1.1):
    try:
        m = model.load()
        out = m(
            prompt,
            max_tokens=min(max_tokens, 8192),
            temperature=temperature,
            top_p=top_p,
            repeat_penalty=repeat_penalty,
            stop=["<|im_end|>", "<|endoftext|>"],
            echo=False,
        )
        return out["choices"][0]["text"].strip()
    except Exception as e:
        logger.error(f"Generate failed: {e}")
        return f"⚠️ Error: {str(e)}"

@spaces.GPU(duration=300)
def chat_respond(message, history, max_tokens=1024, temperature=0.7, top_p=0.9):
    try:
        m = model.load()
        prompt = "<|im_start|>system\nYou are a helpful assistant with 32K context.<|im_end|>\n"
        for h in history:
            prompt += f"<|im_start|>user\n{h[0]}<|im_end|>\n<|im_start|>assistant\n{h[1]}<|im_end|>\n"
        prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
        out = m(
            prompt,
            max_tokens=min(max_tokens, 8192),
            temperature=temperature,
            top_p=top_p,
            stop=["<|im_end|>", "<|endoftext|>"],
            echo=False,
        )
        return out["choices"][0]["text"].strip()
    except Exception as e:
        logger.error(f"Chat failed: {e}")
        return f"⚠️ Error: {str(e)}"

@spaces.GPU(duration=300)
def analyze_image(img, prompt_text):
    if img is None:
        return "Please upload an image first."
    try:
        m = model.load()
        if not model._has_mmproj:
            return "⚠️ Multimodal projection model not available for this GGUF."
        # Convert image to base64 for multimodal
        import base64
        with open(img, "rb") as f:
            b64 = base64.b64encode(f.read()).decode("utf-8")
        ext = Path(img).suffix.lower().lstrip(".")
        if ext in ("jpg", "jpeg"):
            mime = "image/jpeg"
        else:
            mime = f"image/{ext}"
        data_uri = f"data:{mime};base64,{b64}"
        out = m.create_chat_completion(messages=[{
            "role": "user",
            "content": [
                {"type": "text", "text": prompt_text or "Describe this image in detail."},
                {"type": "image_url", "image_url": {"url": data_uri}},
            ],
        }], max_tokens=512, temperature=0.7)
        return out["choices"][0]["message"]["content"]
    except Exception as e:
        logger.error(f"Image analysis failed: {e}")
        return f"⚠️ Error: {str(e)}"

with gr.Blocks(
    title="Gemma-4-E4B Turbo (ZeroGPU)",
    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="green"),
) as demo:
    gr.Markdown("# 🤖 Gemma-4-E4B-Turbo · ZeroGPU\n### 32K Context · 4-bit Q5_K_P · Multimodal")
    
    with gr.Tabs():
        with gr.Tab("💬 Chat"):
            gr.ChatInterface(
                fn=chat_respond,
                additional_inputs=[
                    gr.Slider(128, 8192, value=1024, step=128, label="Max Tokens"),
                    gr.Slider(0.1, 2.0, value=0.7, step=0.05, label="Temperature"),
                    gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-P"),
                ],
            )
        
        with gr.Tab("✍️ Text Generation"):
            with gr.Row():
                with gr.Column(scale=1):
                    prompt = gr.Textbox(lines=6, label="📝 Prompt")
                    with gr.Row():
                        max_tok = gr.Slider(128, 8192, value=1024, step=128, label="Max Tokens")
                        temp = gr.Slider(0.1, 2.0, value=0.7, step=0.05, label="Temperature")
                    top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-P")
                    submit = gr.Button("🚀 Generate", variant="primary")
                with gr.Column(scale=1):
                    output = gr.Textbox(lines=20, label="📄 Output")
            submit.click(fn=generate, inputs=[prompt, max_tok, temp, top_p], outputs=output)
            prompt.submit(fn=generate, inputs=[prompt, max_tok, temp, top_p], outputs=output)
        
        with gr.Tab("🖼️ Image Analysis"):
            gr.Interface(
                fn=analyze_image,
                inputs=[
                    gr.Image(label="Upload Image", type="filepath"),
                    gr.Textbox(label="Prompt (optional)", lines=2, placeholder="Describe this image in detail."),
                ],
                outputs=gr.Textbox(lines=15, label="Analysis"),
                title=None,
                allow_flagging="never",
            )
    
    gr.Markdown("---\n⚡ **ZeroGPU** | Gemma-4-E4B Q5_K_P | First load downloads model (~3.5GB)")

demo.queue(max_size=10).launch()