Update app.py

app.py

@@ -6,7 +6,7 @@ import logging
 import os
 import sys
 import time
-import spaces
+from huggingface_hub import spaces
 import gradio as gr
 import torch
 from PIL import Image
@@ -34,6 +34,10 @@ logger = logging.getLogger("gradio_web_server")
 LOGDIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "logs")
 os.makedirs(os.path.join(LOGDIR, "serve_images"), exist_ok=True)
 
+# Get default model from environment variable or use a fallback
+DEFAULT_MODEL = os.environ.get("DEFAULT_MODEL", "AIML-TUDA/LlavaGuard-v1.2-7B-OV-hf")
+logger.info(f"Using model: {DEFAULT_MODEL}")
+
 default_taxonomy = policy_v1
 
 
@@ -147,6 +151,7 @@ disable_btn = gr.Button(interactive=False)
 
 
 # Model loading function
+@spaces.GPU
 def load_model(model_path):
     global tokenizer, model, processor, context_len
 
@@ -183,16 +188,6 @@ def load_model(model_path):
         return  # Remove return value to avoid Gradio warnings
 
 
-def get_model_list():
-    models = [
-        'AIML-TUDA/QwenGuard-v1.2-3B',
-        'AIML-TUDA/QwenGuard-v1.2-7B',
-        'AIML-TUDA/LlavaGuard-v1.2-0.5B-OV-hf',
-        'AIML-TUDA/LlavaGuard-v1.2-7B-OV-hf',
-    ]
-    return models
-
-
 def get_conv_log_filename():
     t = datetime.datetime.now()
     name = os.path.join(LOGDIR, f"{t.year}-{t.month:02d}-{t.day:02d}-conv.json")
@@ -206,7 +201,7 @@ def run_inference(prompt, image, temperature=0.2, top_p=0.95, max_tokens=512):
     global model, tokenizer, processor
 
     if model is None or processor is None:
-        return "Model not loaded. Please select a model first."
+        return "Model not loaded. Please wait for model to initialize."
     try:
         # Check if it's a Qwen model
         if isinstance(model, Qwen2_5_VLForConditionalGeneration):
@@ -290,57 +285,43 @@ function() {
 
 def load_demo(url_params, request: gr.Request):
     logger.info(f"load_demo. ip: {request.client.host}. params: {url_params}")
-    models = get_model_list()
-
-    dropdown_update = gr.Dropdown(visible=True)
-    if "model" in url_params:
-        model = url_params["model"]
-        if model in models:
-            dropdown_update = gr.Dropdown(value=model, visible=True)
-            load_model(model)
-
     state = default_conversation.copy()
-    return state, dropdown_update
+    return state
 
 
-def load_demo_refresh_model_list(request: gr.Request):
+def load_demo_refresh(request: gr.Request):
     logger.info(f"load_demo. ip: {request.client.host}")
-    models = get_model_list()
     state = default_conversation.copy()
-    dropdown_update = gr.Dropdown(
-        choices=models,
-        value=models[0] if len(models) > 0 else ""
-    )
-    return state, dropdown_update
+    return state
 
 
-def vote_last_response(state, vote_type, model_selector, request: gr.Request):
+def vote_last_response(state, vote_type, request: gr.Request):
     with open(get_conv_log_filename(), "a") as fout:
         data = {
             "tstamp": round(time.time(), 4),
             "type": vote_type,
-            "model": model_selector,
+            "model": DEFAULT_MODEL,
             "state": state.dict(),
             "ip": request.client.host,
         }
         fout.write(json.dumps(data) + "\n")
 
 
-def upvote_last_response(state, model_selector, request: gr.Request):
+def upvote_last_response(state, request: gr.Request):
     logger.info(f"upvote. ip: {request.client.host}")
-    vote_last_response(state, "upvote", model_selector, request)
+    vote_last_response(state, "upvote", request)
     return ("",) + (disable_btn,) * 3
 
 
-def downvote_last_response(state, model_selector, request: gr.Request):
+def downvote_last_response(state, request: gr.Request):
     logger.info(f"downvote. ip: {request.client.host}")
-    vote_last_response(state, "downvote", model_selector, request)
+    vote_last_response(state, "downvote", request)
     return ("",) + (disable_btn,) * 3
 
 
-def flag_last_response(state, model_selector, request: gr.Request):
+def flag_last_response(state, request: gr.Request):
     logger.info(f"flag. ip: {request.client.host}")
-    vote_last_response(state, "flag", model_selector, request)
+    vote_last_response(state, "flag", request)
     return ("",) + (disable_btn,) * 3
 
 
@@ -390,7 +371,7 @@ def add_text(state, text, image, image_process_mode, request: gr.Request):
     return (state, state.to_gradio_chatbot(), default_taxonomy, None) + (disable_btn,) * 5
 
 
-def llava_bot(state, model_selector, temperature, top_p, max_new_tokens, request: gr.Request):
+def llava_bot(state, temperature, top_p, max_new_tokens, request: gr.Request):
     start_tstamp = time.time()
 
     if state.skip_next:
@@ -410,10 +391,6 @@ def llava_bot(state, model_selector, temperature, top_p, max_new_tokens, request
         yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 5
         return
 
-    # Load model if needed
-    if model is None or model_selector != getattr(model, "_name_or_path", ""):
-        load_model(model_selector)
-
     # Run inference
     output = run_inference(prompt, all_images[0], temperature, top_p, max_new_tokens)
 
@@ -434,7 +411,7 @@ def llava_bot(state, model_selector, temperature, top_p, max_new_tokens, request
             data = {
                 "tstamp": round(finish_tstamp, 4),
                 "type": "chat",
-                "model": model_selector,
+                "model": DEFAULT_MODEL,
                 "start": round(start_tstamp, 4),
                 "finish": round(finish_tstamp, 4),
                 "state": state.dict(),
@@ -477,8 +454,6 @@ block_css = """
 
 
 def build_demo(embed_mode, cur_dir=None, concurrency_count=10):
-    models = get_model_list()
-
     with gr.Blocks(title="LlavaGuard", theme=gr.themes.Default(), css=block_css) as demo:
         state = gr.State()
 
@@ -487,13 +462,7 @@ def build_demo(embed_mode, cur_dir=None, concurrency_count=10):
 
         with gr.Row():
             with gr.Column(scale=3):
-                with gr.Row(elem_id="model_selector_row"):
-                    model_selector = gr.Dropdown(
-                        choices=models,
-                        value=models[0] if len(models) > 0 else "",
-                        interactive=True,
-                        show_label=False,
-                        container=False)
+                # Model selector removed
 
                 imagebox = gr.Image(type="pil", label="Image", container=False)
                 image_process_mode = gr.Radio(
@@ -559,35 +528,29 @@ def build_demo(embed_mode, cur_dir=None, concurrency_count=10):
 
         upvote_btn.click(
             upvote_last_response,
-            [state, model_selector],
+            [state],
             [textbox, upvote_btn, downvote_btn, flag_btn]
         )
 
         downvote_btn.click(
             downvote_last_response,
-            [state, model_selector],
+            [state],
             [textbox, upvote_btn, downvote_btn, flag_btn]
         )
 
         flag_btn.click(
             flag_last_response,
-            [state, model_selector],
+            [state],
             [textbox, upvote_btn, downvote_btn, flag_btn]
         )
 
-        model_selector.change(
-            load_model,
-            [model_selector],
-            None
-        )
-
         regenerate_btn.click(
             regenerate,
             [state, image_process_mode],
             [state, chatbot, textbox, imagebox] + btn_list
         ).then(
             llava_bot,
-            [state, model_selector, temperature, top_p, max_output_tokens],
+            [state, temperature, top_p, max_output_tokens],
             [state, chatbot] + btn_list,
             concurrency_limit=concurrency_count
         )
@@ -606,7 +569,7 @@ def build_demo(embed_mode, cur_dir=None, concurrency_count=10):
             queue=False
         ).then(
             llava_bot,
-            [state, model_selector, temperature, top_p, max_output_tokens],
+            [state, temperature, top_p, max_output_tokens],
             [state, chatbot] + btn_list,
             concurrency_limit=concurrency_count
         )
@@ -617,15 +580,15 @@ def build_demo(embed_mode, cur_dir=None, concurrency_count=10):
             [state, chatbot, textbox, imagebox] + btn_list
         ).then(
             llava_bot,
-            [state, model_selector, temperature, top_p, max_output_tokens],
+            [state, temperature, top_p, max_output_tokens],
             [state, chatbot] + btn_list,
             concurrency_limit=concurrency_count
         )
 
         demo.load(
-            load_demo_refresh_model_list,
+            load_demo_refresh,
             None,
-            [state, model_selector],
+            [state],
             queue=False
         )
 
@@ -658,6 +621,8 @@ if api_key:
     login(token=api_key)
     logger.info("Logged in to Hugging Face Hub")
 
+# Load model at startup
+load_model(DEFAULT_MODEL)
 
 demo = build_demo(embed_mode=args.embed, cur_dir='./', concurrency_count=args.concurrency_count)
 demo.queue(
@@ -667,4 +632,4 @@ demo.queue(
     server_name=args.host,
     server_port=args.port,
     share=args.share
-)
+)