Update app.py

app.py

@@ -1,130 +1,93 @@
-# app.py
-# TinyLLaVA Gradio app tailored for CPU-limited Hugging Face Spaces.
-# This file delays heavy model load until the first inference request to reduce build-time memory use.
-
-import io
-import requests
-import traceback
-from typing import Optional
+# app.py - Gradio UI using vendored tinyllava
+import io, requests, traceback
 from PIL import Image
 import gradio as gr
 import torch
 
-# Lazy imports for tinyllava and transformers; we import them only when needed
-_model_loaded = False
-model = None
-tokenizer = None
-image_processor = None
-context_len = None
-device = "cpu"  # Spaces are CPU-only
+from tinyllava.model import load_pretrained_model
+from tinyllava.mm_utils import get_model_name_from_path
 
-# Choose a TinyLLaVA model suitable for CPU (1.5B recommended)
-MODEL_PATH = "bczhou/TinyLLaVA-1.5B"  # recommended for Spaces CPU
+# Use CPU-friendly TinyLLaVA model recommended for Spaces
+MODEL_PATH = "bczhou/TinyLLaVA-1.5B"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 DEFAULT_MAX_TOKENS = 128
-DEFAULT_CONV_MODE = "v1"
 
-def lazy_load_model():
-    global _model_loaded, model, tokenizer, image_processor, context_len
-    if _model_loaded:
-        return
-    try:
-        # Import here after torch is installed by start.py (or already present)
-        from tinyllava.model.builder import load_pretrained_model
-        from tinyllava.mm_utils import get_model_name_from_path
-        from transformers import logging as hf_logging
-        hf_logging.set_verbosity_error()
-    except Exception as e:
-        raise RuntimeError(f"Failed to import TinyLLaVA or transformers: {e}")
+# Lazy load
+_model = None
+_tokenizer = None
+_image_processor = None
+_context_len = None
 
+def lazy_load():
+    global _model, _tokenizer, _image_processor, _context_len
+    if _model is not None:
+        return
     model_name = get_model_name_from_path(MODEL_PATH)
-    tokenizer, model, image_processor, context_len = load_pretrained_model(
-        model_path=MODEL_PATH, model_base=None, model_name=model_name
-    )
-    model.to(device)
-    model.eval()
-    _model_loaded = True
+    _tokenizer, _model, _image_processor, _context_len = load_pretrained_model(MODEL_PATH, model_name=model_name)
+    _model.to(DEVICE)
+    _model.eval()
 
-def load_image_from_url(url: str) -> Image.Image:
+def load_image_from_url(url: str):
     resp = requests.get(url, timeout=10)
     resp.raise_for_status()
     return Image.open(io.BytesIO(resp.content)).convert("RGB")
 
-def _prepare_inputs(prompt: str, image: Image.Image):
+def prepare_inputs(prompt: str, image: Image.Image):
     prompt_text = f"USER: <image>\n{(prompt or '').strip()}\nASSISTANT:"
-    text_inputs = tokenizer(prompt_text, return_tensors="pt").to(device)
-    img_inputs = image_processor(images=image, return_tensors="pt").to(device)
-    inputs = {**text_inputs, **img_inputs}
-    return inputs, prompt_text
+    inputs = _tokenizer(prompt_text, return_tensors="pt").to(DEVICE)
+    if _image_processor is not None:
+        img_inputs = _image_processor(images=image, return_tensors="pt").to(DEVICE)
+        inputs.update(img_inputs)
+    return inputs
 
-def generate_text(prompt: str, uploaded_image: Optional[Image.Image], image_url: str,
-                  max_new_tokens: int = DEFAULT_MAX_TOKENS, conv_mode: str = DEFAULT_CONV_MODE):
+def generate_text(prompt, upload, url, max_new_tokens=DEFAULT_MAX_TOKENS):
     try:
-        # Ensure model is loaded on first call
-        if not _model_loaded:
-            lazy_load_model()
-
-        if uploaded_image is None and image_url:
-            image = load_image_from_url(image_url)
-        elif uploaded_image is not None:
-            image = uploaded_image
+        lazy_load()
+        if upload is None and url:
+            image = load_image_from_url(url)
+        elif upload is not None:
+            image = upload
         else:
-            return "No image provided. Upload an image or provide an image URL."
-
-        inputs, prompt_text = _prepare_inputs(prompt, image)
-
-        gen_kwargs = {
-            "max_new_tokens": int(max_new_tokens),
-            "num_beams": 1,
-            "temperature": 0.0,
-        }
-
-        outputs = model.generate(**inputs, **gen_kwargs)
-
-        # Decode outputs
-        if isinstance(outputs, torch.Tensor):
-            decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
-        elif isinstance(outputs, (list, tuple)):
-            decoded = tokenizer.batch_decode(outputs[0], skip_special_tokens=True)[0]
+            return "No image provided."
+
+        inputs = prepare_inputs(prompt, image)
+        gen = _model.generate(**inputs, max_new_tokens=int(max_new_tokens), num_beams=1, temperature=0.0)
+        if isinstance(gen, torch.Tensor):
+            decoded = _tokenizer.batch_decode(gen, skip_special_tokens=True)[0]
+        elif isinstance(gen, (list, tuple)):
+            decoded = _tokenizer.batch_decode(gen[0], skip_special_tokens=True)[0]
         else:
-            decoded = str(outputs)
-
+            decoded = str(gen)
         if "ASSISTANT:" in decoded:
-            reply = decoded.split("ASSISTANT:")[-1].strip()
-        else:
-            reply = decoded.strip()
-        return reply
+            return decoded.split("ASSISTANT:")[-1].strip()
+        return decoded.strip()
     except Exception as e:
-        tb = traceback.format_exc()
-        return f"Inference error: {e}\n\nTraceback:\n{tb}"
+        return f"Inference error: {e}\n\n{traceback.format_exc()}"
 
-# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# TinyLLaVA (CPU) — Hugging Face Spaces Demo")
+    gr.Markdown("TinyLLaVA (vendored loader) — CPU-friendly")
     with gr.Row():
         with gr.Column(scale=2):
-            prompt_input = gr.Textbox(label="Prompt (optional)", placeholder="Ask about the image...")
-            upload = gr.Image(label="Upload Image (preferred)", type="pil")
-            url = gr.Textbox(label="Image URL (used if upload empty)", placeholder="https://...")
-            max_tokens = gr.Slider(minimum=32, maximum=512, step=32, value=DEFAULT_MAX_TOKENS, label="Max new tokens")
-            run_btn = gr.Button("Generate")
+            prompt = gr.Textbox(label="Prompt (optional)")
+            upload = gr.Image(label="Upload Image", type="pil")
+            url = gr.Textbox(label="Image URL")
+            max_tokens = gr.Slider(32, 512, value=DEFAULT_MAX_TOKENS, step=32, label="Max new tokens")
+            btn = gr.Button("Generate")
         with gr.Column(scale=1):
-            preview = gr.Image(label="Image preview", type="pil")
-            out = gr.Textbox(label="Generated Text", lines=8)
-
-    def update_preview(uploaded, url_text):
-        if uploaded is not None:
-            return uploaded
-        if url_text:
+            preview = gr.Image(label="Preview", type="pil")
+            out = gr.Textbox(label="Output", lines=8)
+    def update_preview(u, s):
+        if u is not None:
+            return u
+        if s:
             try:
-                return load_image_from_url(url_text)
-            except Exception:
+                return load_image_from_url(s)
+            except:
                 return None
         return None
-
-    upload.change(fn=update_preview, inputs=[upload, url], outputs=preview)
-    url.change(fn=update_preview, inputs=[upload, url], outputs=preview)
-
-    run_btn.click(fn=generate_text, inputs=[prompt_input, upload, url, max_tokens], outputs=out)
+    upload.change(update_preview, [upload, url], preview)
+    url.change(update_preview, [upload, url], preview)
+    btn.click(generate_text, [prompt, upload, url, max_tokens], out)
 
 if __name__ == "__main__":
     demo.launch()