Update app.py

app.py

@@ -12,7 +12,7 @@ import torch
 from PIL import Image
 from transformers import (
     AutoProcessor,
-    AutoTokenizer, 
+    AutoTokenizer,
     Qwen2_5_VLForConditionalGeneration,
     LlavaOnevisionForConditionalGeneration
 )
@@ -74,16 +74,16 @@ class SimpleConversation:
     def to_gradio_chatbot(self):
         if not self.messages:
             return []
-            
+
         ret = []
         for msg in self.messages:
             prompt = msg[0]
             if isinstance(prompt, tuple) and len(prompt) > 0:
                 prompt = prompt[0]
-                
+
             if prompt and isinstance(prompt, str) and "<image>" in prompt:
                 prompt = prompt.replace("<image>", "")
-                
+
             ret.append([prompt, msg[1]])
         return ret
 
@@ -123,6 +123,7 @@ class SimpleConversation:
         new_conv.messages = self.messages.copy() if self.messages else []
         return new_conv
 
+
 default_conversation = SimpleConversation()
 
 # Model and processor storage
@@ -131,55 +132,56 @@ model = None
 processor = None
 context_len = 8048
 
+
 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
 
+
 # UI component states
 no_change_btn = gr.Button()
 enable_btn = gr.Button(interactive=True)
 disable_btn = gr.Button(interactive=False)
 
+
 # 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,
+                model_path,
                 torch_dtype="auto",
                 device_map="auto"
             )
             processor = AutoProcessor.from_pretrained(model_path)
             tokenizer = processor.tokenizer
-            
+
         # Otherwise assume it's a LlavaGuard model
         else:
             model = LlavaOnevisionForConditionalGeneration.from_pretrained(
                 model_path,
-                # torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
                 torch_dtype="auto",
                 device_map="auto",
                 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", 8048)
-        logger.info(f"Model {model_path} loaded successfully to device: {model.device}")
-        model = model.to("cuda")
-        return True
+        logger.info(f"Model {model_path} loaded successfully")
+        return  # Remove return value to avoid Gradio warnings
 
     except Exception as e:
         logger.error(f"Error loading model {model_path}: {str(e)}")
-        return False
+        return  # Remove return value to avoid Gradio warnings
+
 
 def get_model_list():
     models = [
@@ -190,17 +192,19 @@ def get_model_list():
     ]
     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
 @spaces.GPU
 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:
@@ -227,7 +231,7 @@ def run_inference(prompt, image, temperature=0.2, top_p=0.95, max_tokens=512):
                 return_tensors="pt",
             )
 
-            
+
         # Otherwise assume it's a LlavaGuard model
         else:
             conversation = [
@@ -272,6 +276,7 @@ def run_inference(prompt, image, temperature=0.2, top_p=0.95, max_tokens=512):
         logger.error(error_msg)
         return f"Error processing image. Please try again."
 
+
 # Gradio UI functions
 get_window_url_params = """
 function() {
@@ -282,10 +287,11 @@ 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"]
@@ -296,6 +302,7 @@ def load_demo(url_params, request: gr.Request):
     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()
@@ -306,6 +313,7 @@ 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 = {
@@ -317,21 +325,25 @@ 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}")
     if state.messages and len(state.messages) > 0:
@@ -344,15 +356,17 @@ def regenerate(state, image_process_mode, request: gr.Request):
                 if len(prev_human_msg[0]) >= 3:
                     new_msg[0] = (prev_human_msg[0][0], prev_human_msg[0][1], image_process_mode)
                 state.messages[-2] = new_msg
-    
+
     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:
@@ -360,24 +374,25 @@ def add_text(state, text, image, image_process_mode, request: gr.Request):
         return (state, state.to_gradio_chatbot(), "", None) + (no_change_btn,) * 5
 
     text = wrap_taxonomy(text)
-    
+
     # Reset conversation for new image-based query
     if image is not None:
         state = default_conversation.copy()
-        
+
     # Set new prompt with image
     prompt = text
     if image is not None:
         prompt = (text, image, image_process_mode)
-    
+
     state.set_prompt(prompt=prompt, image=image)
     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()
-    
+
     if state.skip_next:
         # This generate call is skipped due to invalid inputs
         yield (state, state.to_gradio_chatbot()) + (no_change_btn,) * 5
@@ -386,7 +401,7 @@ def llava_bot(state, model_selector, temperature, top_p, max_new_tokens, request
     # Get the prompt and images
     prompt = state.get_prompt()
     all_images = state.get_image(return_pil=True)
-    
+
     if not all_images:
         if not state.messages:
             state.messages = [["Error: No image provided", None]]
@@ -394,14 +409,14 @@ def llava_bot(state, model_selector, temperature, top_p, max_new_tokens, request
             state.messages[-1][-1] = "Error: No image provided"
         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)
-    
+
     # Update the response in the conversation state
     if not state.messages:
         state.messages = [[prompt, output]]
@@ -430,6 +445,7 @@ def llava_bot(state, model_selector, temperature, top_p, max_new_tokens, request
     except Exception as e:
         logger.error(f"Error writing log: {str(e)}")
 
+
 # UI Components
 title_markdown = """
 # LLAVAGUARD: VLM-based Safeguard for Vision Dataset Curation and Safety Assessment
@@ -459,6 +475,7 @@ block_css = """
 }
 """
 
+
 def build_demo(embed_mode, cur_dir=None, concurrency_count=10):
     models = get_model_list()
 
@@ -486,17 +503,18 @@ 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) if os.path.exists(f"{cur_dir}/examples/image{i}.png")
+                    [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")
+                                            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.Accordion("Safety Risk Taxonomy", open=False):
                     taxonomy_textbox = gr.Textbox(
@@ -538,25 +556,25 @@ 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_model,
             [model_selector],
@@ -626,38 +644,39 @@ if __name__ == "__main__":
 
     # 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")
     if api_key:
         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=args.embed, cur_dir='./', concurrency_count=args.concurrency_count)
-        demo.queue(
-            status_update_rate=10,
-            api_open=False
-        ).launch(
-            server_name=args.host,
-            server_port=args.port,
-            share=args.share
-        )
-    except Exception as e:
-        logger.error(f"Error launching demo: {e}")
-        sys.exit(1)
 
+    # Launch Gradio app in a subprocess to avoid CUDA initialization in the main process
+    from torch.multiprocessing import Process
+
+    def launch_demo():
+        try:
+            demo = build_demo(embed_mode=args.embed, cur_dir='./', concurrency_count=args.concurrency_count)
+            demo.queue(
+                status_update_rate=10,
+                api_open=False
+            ).launch(
+                server_name=args.host,
+                server_port=args.port,
+                share=args.share
+            )
+        except Exception as e:
+            logger.error(f"Error launching demo: {e}")
+            sys.exit(1)
+
+    p = Process(target=launch_demo)
+    p.start()
+    p.join()