Update app.py

app.py

@@ -2,120 +2,299 @@ import argparse
 import datetime
 import hashlib
 import json
+import logging
 import os
 import sys
 import time
-import warnings
 
 import gradio as gr
-import spaces
 import torch
-
-from builder import load_pretrained_model
-from llava.constants import IMAGE_TOKEN_INDEX
-from llava.constants import LOGDIR
-from llava.conversation import (default_conversation, conv_templates)
-from llava.mm_utils import KeywordsStoppingCriteria, tokenizer_image_token
-from llava.utils import (build_logger, violates_moderation, moderation_msg)
-from taxonomy import wrap_taxonomy, default_taxonomy
-
-
+from PIL import Image
+from transformers import (
+    AutoModelForCausalLM, 
+    AutoProcessor, 
+    AutoTokenizer, 
+    Qwen2_5_VLForConditionalGeneration
+)
+
+from taxonomy import policy_v1
+
+# Set up logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.FileHandler("gradio_web_server.log"),
+        logging.StreamHandler()
+    ]
+)
+logger = logging.getLogger("gradio_web_server")
+
+# Constants
+LOGDIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "logs")
+os.makedirs(os.path.join(LOGDIR, "serve_images"), exist_ok=True)
+
+default_taxonomy = policy_v1
+
+class Conversation:
+    def __init__(self):
+        self.messages = []
+        self.roles = ["user", "assistant"]
+        self.offset = 0
+        self.skip_next = False
+
+    def append_message(self, role, message):
+        self.messages.append([role, message])
+
+    def to_gradio_chatbot(self):
+        ret = []
+        for role, message in self.messages:
+            if message is None:
+                continue
+            if role == self.roles[0]:
+                if isinstance(message, tuple):
+                    ret.append([self.render_user_message(message[0]), None])
+                else:
+                    ret.append([self.render_user_message(message), None])
+            elif role == self.roles[1]:
+                if ret[-1][1] is None:
+                    ret[-1][1] = message
+                else:
+                    ret.append([None, message])
+            else:
+                raise ValueError(f"Invalid role: {role}")
+        return ret
+    
+    def render_user_message(self, message):
+        if "<image>" in message:
+            return message.replace("<image>", "")
+        return message
+    
+    def dict(self):
+        return {
+            "messages": self.messages,
+            "roles": self.roles,
+            "offset": self.offset,
+            "skip_next": self.skip_next,
+        }
+    
+    def get_prompt(self):
+        prompt = ""
+        for role, message in self.messages:
+            if message is None:
+                continue
+            if isinstance(message, tuple):
+                message = message[0]
+            if role == self.roles[0]:
+                prompt += f"USER: {message}\n"
+            else:
+                prompt += f"ASSISTANT: {message}\n"
+        return prompt + "ASSISTANT: "
+    
+    def get_images(self, return_pil=False):
+        images = []
+        for role, message in self.messages:
+            if isinstance(message, tuple) and len(message) > 1:
+                if isinstance(message[1], Image.Image):
+                    images.append(message[1] if return_pil else message[1])
+        return images
+    
+    def copy(self):
+        new_conv = Conversation()
+        new_conv.messages = self.messages.copy()
+        new_conv.roles = self.roles.copy()
+        new_conv.offset = self.offset
+        new_conv.skip_next = self.skip_next
+        return new_conv
+
+default_conversation = Conversation()
+
+# Model and processor storage
+tokenizer = None
+model = None
+processor = None
+context_len = 2048
+
+# Helper functions
 def clear_conv(conv):
     conv.messages = []
     return conv
 
+def wrap_taxonomy(text):
+    """Wraps user input with taxonomy if not already present"""
+    if policy_v1 not in text:
+        return policy_v1 + "\n\n" + text
+    return text
 
-logger = build_logger("gradio_web_server", "gradio_web_server.log")
-
-headers = {"User-Agent": "LLaVA Client"}
-
+# UI component states
 no_change_btn = gr.Button()
 enable_btn = gr.Button(interactive=True)
 disable_btn = gr.Button(interactive=False)
 
-priority = {
-    "LlavaGuard-7B": "aaaaaaa",
-    "LlavaGuard-13B": "aaaaaab",
-    "LlavaGuard-34B": "aaaaaac",
-}
-
-
-@spaces.GPU
-def run_llava(prompt, pil_image, temperature, top_p, max_new_tokens):
-    image_size = pil_image.size
-    image_tensor = image_processor.preprocess(pil_image, return_tensors='pt')['pixel_values'].half().cuda()
-    # image_tensor = image_tensor.to(model.device, dtype=torch.float16)
-    input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt")
-    input_ids = input_ids.unsqueeze(0).cuda()
-    with torch.inference_mode():
-        output_ids = model.generate(
-            input_ids,
-            images=image_tensor,
-            image_sizes=[image_size],
-            do_sample=True,
-            temperature=temperature,
-            top_p=top_p,
-            top_k=50,
-            num_beams=2,
-            max_new_tokens=max_new_tokens,
-            use_cache=True,
-            stopping_criteria=[KeywordsStoppingCriteria(['}'], tokenizer, input_ids)]
-        )
-    outputs = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
-
-    return outputs[0].strip()
-
-
-def load_selected_model(model_path):
-    model_name = model_path.split("/")[-1]
-    global tokenizer, model, image_processor, context_len
-    with warnings.catch_warnings(record=True) as w:
-        warnings.simplefilter("always")
-        tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, None, model_name)
-    for warning in w:
-        if "vision" not in str(warning.message).lower():
-            print(warning.message)
-    model.config.tokenizer_model_max_length = 2048 * 2
+# Model loading function
+def load_model(model_path):
+    global tokenizer, model, processor, context_len
+    
+    logger.info(f"Loading model: {model_path}")
+    
+    try:
+        # Check if it's a Qwen model
+        if "qwenguard" in model_path.lower():
+            model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+                model_path, 
+                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                device_map="auto" if torch.cuda.is_available() else None
+            )
+            processor = AutoProcessor.from_pretrained(model_path)
+            tokenizer = processor.tokenizer
+            
+        # Otherwise assume it's a LlavaGuard model
+        else:
+            model = AutoModelForCausalLM.from_pretrained(
+                model_path,
+                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                device_map="auto" if torch.cuda.is_available() else None,
+                trust_remote_code=True
+            )
+            tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+            processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
+        
+        context_len = getattr(model.config, "max_position_embeddings", 2048)
+        logger.info(f"Model {model_path} loaded successfully")
+        return True
+    
+    except Exception as e:
+        logger.error(f"Error loading model {model_path}: {str(e)}")
+        return False
 
+def get_model_list():
+    models = [
+        'AIML-TUDA/LlavaGuard-v1.2-7B-OV-hf',
+        'AIML-TUDA/LlavaGuard-v1.2-0.5B-OV-hf',
+        'AIML-TUDA/QwenGuard-v1.2-7B',
+        'AIML-TUDA/QwenGuard-v1.2-3B'
+    ]
+    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")
+    os.makedirs(os.path.dirname(name), exist_ok=True)
     return name
 
+# Inference function
+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."
+    
+    try:
+        # Check if it's a Qwen model
+        if isinstance(model, Qwen2_5_VLForConditionalGeneration):
+            # Format for Qwen models
+            messages = [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "image", "image": image},
+                        {"type": "text", "text": prompt}
+                    ]
+                }
+            ]
+            
+            # Process input
+            text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+            inputs = processor(
+                text=[text],
+                images=[image], 
+                padding=True, 
+                return_tensors="pt"
+            )
+            
+            # Move to GPU if available
+            if torch.cuda.is_available():
+                inputs = {k: v.to("cuda") for k, v in inputs.items()}
+            
+            # Generate
+            with torch.no_grad():
+                generated_ids = model.generate(
+                    **inputs,
+                    do_sample=temperature > 0,
+                    temperature=temperature,
+                    top_p=top_p,
+                    max_new_tokens=max_tokens,
+                )
+                
+            # Decode
+            generated_ids_trimmed = [
+                out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+            ]
+            response = processor.batch_decode(
+                generated_ids_trimmed, 
+                skip_special_tokens=True,
+                clean_up_tokenization_spaces=False
+            )[0]
+            
+        # Otherwise assume it's a LlavaGuard model
+        else:
+            # Process input for LlavaGuard models
+            inputs = processor(
+                prompt, 
+                images=image, 
+                return_tensors="pt"
+            )
+            
+            # Move to GPU if available
+            if torch.cuda.is_available():
+                inputs = {k: v.to("cuda") for k, v in inputs.items()}
+            
+            # Generate
+            with torch.no_grad():
+                generated_ids = model.generate(
+                    **inputs,
+                    do_sample=temperature > 0,
+                    temperature=temperature,
+                    top_p=top_p,
+                    max_new_tokens=max_tokens,
+                )
+            
+            # Decode
+            response = tokenizer.batch_decode(
+                generated_ids[:, inputs.input_ids.shape[1]:], 
+                skip_special_tokens=True
+            )[0]
+        
+        return response.strip()
+        
+    except Exception as e:
+        logger.error(f"Error during inference: {str(e)}")
+        return f"Error during inference: {str(e)}"
 
-def get_model_list():
-    models = [
-                 'AIML-TUDA/LlavaGuard-7B',
-                 'AIML-TUDA/LlavaGuard-v1.1-7B-hf',
-                 'AIML-TUDA/LlavaGuard-13B',
-                 'AIML-TUDA/LlavaGuard-v1.1-13B-hf']
-    return models
-
-
+# Gradio UI functions
 get_window_url_params = """
 function() {
     const params = new URLSearchParams(window.location.search);
     url_params = Object.fromEntries(params);
     console.log(url_params);
     return url_params;
-    }
+}
 """
 
-
 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
 
-
 def load_demo_refresh_model_list(request: gr.Request):
     logger.info(f"load_demo. ip: {request.client.host}")
     models = get_model_list()
@@ -126,7 +305,6 @@ def load_demo_refresh_model_list(request: gr.Request):
     )
     return state, dropdown_update
 
-
 def vote_last_response(state, vote_type, model_selector, request: gr.Request):
     with open(get_conv_log_filename(), "a") as fout:
         data = {
@@ -138,25 +316,21 @@ def vote_last_response(state, vote_type, model_selector, request: gr.Request):
         }
         fout.write(json.dumps(data) + "\n")
 
-
 def upvote_last_response(state, model_selector, request: gr.Request):
     logger.info(f"upvote. ip: {request.client.host}")
     vote_last_response(state, "upvote", model_selector, request)
     return ("",) + (disable_btn,) * 3
 
-
 def downvote_last_response(state, model_selector, request: gr.Request):
     logger.info(f"downvote. ip: {request.client.host}")
     vote_last_response(state, "downvote", model_selector, request)
     return ("",) + (disable_btn,) * 3
 
-
 def flag_last_response(state, model_selector, request: gr.Request):
     logger.info(f"flag. ip: {request.client.host}")
     vote_last_response(state, "flag", model_selector, request)
     return ("",) + (disable_btn,) * 3
 
-
 def regenerate(state, image_process_mode, request: gr.Request):
     logger.info(f"regenerate. ip: {request.client.host}")
     state.messages[-1][-1] = None
@@ -166,30 +340,20 @@ def regenerate(state, image_process_mode, request: gr.Request):
     state.skip_next = False
     return (state, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 5
 
-
 def clear_history(request: gr.Request):
     logger.info(f"clear_history. ip: {request.client.host}")
     state = default_conversation.copy()
     return (state, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 5
 
-
 def add_text(state, text, image, image_process_mode, request: gr.Request):
     logger.info(f"add_text. ip: {request.client.host}. len: {len(text)}")
     if len(text) <= 0 or image is None:
         state.skip_next = True
         return (state, state.to_gradio_chatbot(), "", None) + (no_change_btn,) * 5
-    if args.moderate:
-        flagged = violates_moderation(text)
-        if flagged:
-            state.skip_next = True
-            return (state, state.to_gradio_chatbot(), moderation_msg, None) + (
-                no_change_btn,) * 5
 
     text = wrap_taxonomy(text)
     if image is not None:
-        text = text  # Hard cut-off for images
         if '<image>' not in text:
-            # text = '<Image><image></Image>' + text
             text = text + '\n<image>'
         text = (text, image, image_process_mode)
         state = default_conversation.copy()
@@ -199,83 +363,50 @@ def add_text(state, text, image, image_process_mode, request: gr.Request):
     state.skip_next = False
     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):
     start_tstamp = time.time()
-    model_name = model_selector
-
+    
     if state.skip_next:
         # This generate call is skipped due to invalid inputs
         yield (state, state.to_gradio_chatbot()) + (no_change_btn,) * 5
         return
 
-    if len(state.messages) == state.offset + 2:
-        # First round of conversation
-        if "llava" in model_name.lower():
-            if 'llama-2' in model_name.lower():
-                template_name = "llava_llama_2"
-            elif "mistral" in model_name.lower() or "mixtral" in model_name.lower():
-                if 'orca' in model_name.lower():
-                    template_name = "mistral_orca"
-                elif 'hermes' in model_name.lower():
-                    template_name = "chatml_direct"
-                else:
-                    template_name = "mistral_instruct"
-            elif 'llava-v1.6-34b' in model_name.lower():
-                template_name = "chatml_direct"
-            elif "v1" in model_name.lower():
-                if 'mmtag' in model_name.lower():
-                    template_name = "v1_mmtag"
-                elif 'plain' in model_name.lower() and 'finetune' not in model_name.lower():
-                    template_name = "v1_mmtag"
-                else:
-                    template_name = "llava_v1"
-            elif "mpt" in model_name.lower():
-                template_name = "mpt"
-            else:
-                if 'mmtag' in model_name.lower():
-                    template_name = "v0_mmtag"
-                elif 'plain' in model_name.lower() and 'finetune' not in model_name.lower():
-                    template_name = "v0_mmtag"
-                else:
-                    template_name = "llava_v0"
-        elif "mpt" in model_name:
-            template_name = "mpt_text"
-        elif "llama-2" in model_name:
-            template_name = "llama_2"
-        else:
-            template_name = "vicuna_v1"
-        new_state = conv_templates[template_name].copy()
-        new_state.append_message(new_state.roles[0], state.messages[-2][1])
-        new_state.append_message(new_state.roles[1], None)
-        state = new_state
-
-    # Construct prompt
+    # Get the prompt and images
     prompt = state.get_prompt()
-
     all_images = state.get_images(return_pil=True)
+    
+    if not all_images:
+        state.messages[-1][-1] = "Error: No image provided"
+        yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 5
+        return
+    
+    # Save image for logging
     all_image_hash = [hashlib.md5(image.tobytes()).hexdigest() for image in all_images]
-    for image, hash in zip(all_images, all_image_hash):
+    for image, hash_val in zip(all_images, all_image_hash):
         t = datetime.datetime.now()
-        filename = os.path.join(LOGDIR, "serve_images", f"{t.year}-{t.month:02d}-{t.day:02d}", f"{hash}.jpg")
+        filename = os.path.join(LOGDIR, "serve_images", f"{t.year}-{t.month:02d}-{t.day:02d}", f"{hash_val}.jpg")
         if not os.path.isfile(filename):
             os.makedirs(os.path.dirname(filename), exist_ok=True)
             image.save(filename)
-
-    output = run_llava(prompt, all_images[0], temperature, top_p, max_new_tokens)
-
+    
+    # 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)
     state.messages[-1][-1] = output
 
     yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 5
 
     finish_tstamp = time.time()
-    logger.info(f"{output}")
+    logger.info(f"Generated response in {finish_tstamp - start_tstamp:.2f}s")
 
     with open(get_conv_log_filename(), "a") as fout:
         data = {
             "tstamp": round(finish_tstamp, 4),
             "type": "chat",
-            "model": model_name,
+            "model": model_selector,
             "start": round(start_tstamp, 4),
             "finish": round(finish_tstamp, 4),
             "state": state.dict(),
@@ -284,41 +415,38 @@ def llava_bot(state, model_selector, temperature, top_p, max_new_tokens, request
         }
         fout.write(json.dumps(data) + "\n")
 
-
-title_markdown = ("""
+# UI Components
+title_markdown = """
 # LLAVAGUARD: VLM-based Safeguard for Vision Dataset Curation and Safety Assessment
 [[Project Page](https://ml-research.github.io/human-centered-genai/projects/llavaguard/index.html)]
 [[Code](https://github.com/ml-research/LlavaGuard)]
 [[Model](https://huggingface.co/collections/AIML-TUDA/llavaguard-665b42e89803408ee8ec1086)]
 [[Dataset](https://huggingface.co/datasets/aiml-tuda/llavaguard)]
 [[LavaGuard](https://arxiv.org/abs/2406.05113)]
-""")
+"""
 
-tos_markdown = ("""
+tos_markdown = """
 ### Terms of use
 By using this service, users are required to agree to the following terms:
 The service is a research preview intended for non-commercial use only. It only provides limited safety measures and may generate offensive content. It must not be used for any illegal, harmful, violent, racist, or sexual purposes. The service may collect user dialogue data for future research.
 Please click the "Flag" button if you get any inappropriate answer! We will collect those to keep improving our moderator.
 For an optimal experience, please use desktop computers for this demo, as mobile devices may compromise its quality.
-""")
+"""
 
-learn_more_markdown = ("""
+learn_more_markdown = """
 ### License
 The service is a research preview intended for non-commercial use only, subject to the model [License](https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md) of LLaMA, [Terms of Use](https://openai.com/policies/terms-of-use) of the data generated by OpenAI, and [Privacy Practices](https://chrome.google.com/webstore/detail/sharegpt-share-your-chatg/daiacboceoaocpibfodeljbdfacokfjb) of ShareGPT. Please contact us if you find any potential violation.
-""")
+"""
 
 block_css = """
-
 #buttons button {
     min-width: min(120px,100%);
 }
-
 """
 
-taxonomies = ["Default", "Modified w/ O1 non-violating", "Default message 3"]
-
-
 def build_demo(embed_mode, cur_dir=None, concurrency_count=10):
+    models = get_model_list()
+    
     with gr.Accordion("Safety Risk Taxonomy", open=False) as accordion:
         textbox = gr.Textbox(
             label="Safety Risk Taxonomy",
@@ -327,6 +455,7 @@ def build_demo(embed_mode, cur_dir=None, concurrency_count=10):
             container=True,
             value=default_taxonomy,
             lines=50)
+            
     with gr.Blocks(title="LlavaGuard", theme=gr.themes.Default(), css=block_css) as demo:
         state = gr.State()
 
@@ -351,16 +480,17 @@ def build_demo(embed_mode, cur_dir=None, concurrency_count=10):
 
                 if cur_dir is None:
                     cur_dir = os.path.dirname(os.path.abspath(__file__))
+                    
                 gr.Examples(examples=[
-                    [f"{cur_dir}/examples/image{i}.png"] for i in range(1, 6)
+                    [f"{cur_dir}/examples/image{i}.png"] for i in range(1, 6) if os.path.exists(f"{cur_dir}/examples/image{i}.png")
                 ], inputs=imagebox)
 
                 with gr.Accordion("Parameters", open=False) as parameter_row:
                     temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.2, step=0.1, interactive=True,
-                                            label="Temperature", )
-                    top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.95, step=0.1, interactive=True, label="Top P", )
+                                           label="Temperature")
+                    top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.95, step=0.1, interactive=True, label="Top P")
                     max_output_tokens = gr.Slider(minimum=0, maximum=1024, value=512, step=64, interactive=True,
-                                                  label="Max output tokens", )
+                                                 label="Max output tokens")
 
             with gr.Column(scale=8):
                 chatbot = gr.Chatbot(
@@ -378,7 +508,6 @@ def build_demo(embed_mode, cur_dir=None, concurrency_count=10):
                     upvote_btn = gr.Button(value="👍  Upvote", interactive=False)
                     downvote_btn = gr.Button(value="👎  Downvote", interactive=False)
                     flag_btn = gr.Button(value="⚠️  Flag", interactive=False)
-                    # stop_btn = gr.Button(value="⏹️  Stop Generation", interactive=False)
                     regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False)
                     clear_btn = gr.Button(value="🗑️  Clear", interactive=False)
 
@@ -389,26 +518,30 @@ def build_demo(embed_mode, cur_dir=None, concurrency_count=10):
 
         # Register listeners
         btn_list = [upvote_btn, downvote_btn, flag_btn, regenerate_btn, clear_btn]
+        
         upvote_btn.click(
             upvote_last_response,
             [state, model_selector],
             [textbox, upvote_btn, downvote_btn, flag_btn]
         )
+        
         downvote_btn.click(
             downvote_last_response,
             [state, model_selector],
             [textbox, upvote_btn, downvote_btn, flag_btn]
         )
+        
         flag_btn.click(
             flag_last_response,
             [state, model_selector],
             [textbox, upvote_btn, downvote_btn, flag_btn]
         )
-
-        # model_selector.change(
-        #     load_selected_model,
-        #     [model_selector],
-        # )
+        
+        model_selector.change(
+            load_model,
+            [model_selector],
+            None
+        )
 
         regenerate_btn.click(
             regenerate,
@@ -451,22 +584,12 @@ def build_demo(embed_mode, cur_dir=None, concurrency_count=10):
             concurrency_limit=concurrency_count
         )
 
-        if args.model_list_mode == "once":
-            demo.load(
-                load_demo,
-                [url_params],
-                [state, model_selector],
-                js=get_window_url_params
-            )
-        elif args.model_list_mode == "reload":
-            demo.load(
-                load_demo_refresh_model_list,
-                None,
-                [state, model_selector],
-                queue=False
-            )
-        else:
-            raise ValueError(f"Unknown model list mode: {args.model_list_mode}")
+        demo.load(
+            load_demo_refresh_model_list,
+            None,
+            [state, model_selector],
+            queue=False
+        )
 
     return demo
 
@@ -475,51 +598,37 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--host", type=str, default="0.0.0.0")
     parser.add_argument("--port", type=int)
-    parser.add_argument("--controller-url", type=str, default="http://localhost:10000")
     parser.add_argument("--concurrency-count", type=int, default=5)
-    parser.add_argument("--model-list-mode", type=str, default="reload", choices=["once", "reload"])
     parser.add_argument("--share", action="store_true")
     parser.add_argument("--moderate", action="store_true")
     parser.add_argument("--embed", action="store_true")
     args = parser.parse_args()
-    models = []
-
-    title_markdown += """
 
-ONLY WORKS WITH GPU!
-
-Set the environment variable `model` to change the model:
-['AIML-TUDA/LlavaGuard-7B'](https://huggingface.co/AIML-TUDA/LlavaGuard-7B),
-['AIML-TUDA/LlavaGuard-13B'](https://huggingface.co/AIML-TUDA/LlavaGuard-13B),
-['AIML-TUDA/LlavaGuard-34B'](https://huggingface.co/AIML-TUDA/LlavaGuard-34B),
-"""
-    print(f"args: {args}")
-    concurrency_count = int(os.getenv("concurrency_count", 5))
+    # Create log directory if it doesn't exist
+    os.makedirs(LOGDIR, exist_ok=True)
+    
+    # GPU Check
+    if torch.cuda.is_available():
+        logger.info(f"CUDA available with {torch.cuda.device_count()} devices")
+    else:
+        logger.warning("CUDA not available! Models will run on CPU which may be very slow.")
+    
+    # Hugging Face token handling
     api_key = os.getenv("token")
-
-
-    models = get_model_list()
-    bits = int(os.getenv("bits", 16))
-    model = os.getenv("model", models[0])
-    available_devices = os.getenv("CUDA_VISIBLE_DEVICES", "0")
-    model_path, model_name = model, model.split("/")[0]
     if api_key:
-        cmd = f"huggingface-cli login --token {api_key} --add-to-git-credential"
-        os.system(cmd)
-    else:
-        if '/workspace' not in sys.path:
-            sys.path.append('/workspace')
-        from llavaguard.hf_utils import set_up_env_and_token
-        api_key = set_up_env_and_token(read=True, write=False)
-        model_path = '/common-repos/LlavaGuard/models/LlavaGuard-v1.1-7b-full/smid_and_crawled_v2_with_augmented_policies/json-v16/llava'
-
-    print(f"Loading model {model_path}")
-    load_selected_model(model_path)
-    model.config.tokenizer_model_max_length = 2048 * 2
-
-    exit_status = 0
+        from huggingface_hub import login
+        login(token=api_key)
+        logger.info("Logged in to Hugging Face Hub")
+    
+    # Load initial model
+    models = get_model_list()
+    model_path = os.getenv("model", models[0])
+    logger.info(f"Initial model selected: {model_path}")
+    load_model(model_path)
+    
+    # Launch Gradio app
     try:
-        demo = build_demo(embed_mode=False, cur_dir='./', concurrency_count=concurrency_count)
+        demo = build_demo(embed_mode=args.embed, cur_dir='./', concurrency_count=args.concurrency_count)
         demo.queue(
             status_update_rate=10,
             api_open=False
@@ -528,9 +637,6 @@ Set the environment variable `model` to change the model:
             server_port=args.port,
             share=args.share
         )
-
     except Exception as e:
-        print(e)
-        exit_status = 1
-    finally:
-        sys.exit(exit_status)
+        logger.error(f"Error launching demo: {e}")
+        sys.exit(1)