Update app.py

app.py

@@ -16,6 +16,7 @@ from transformers import (
     Qwen2_5_VLForConditionalGeneration,
     LlavaOnevisionForConditionalGeneration
 )
+from qwen_vl_utils import process_vision_info
 
 from taxonomy import policy_v1
 
@@ -36,94 +37,77 @@ os.makedirs(os.path.join(LOGDIR, "serve_images"), exist_ok=True)
 
 default_taxonomy = policy_v1
 
-class Conversation:
+
+class SimpleConversation:
     def __init__(self):
-        self.messages = []
-        self.roles = ["user", "assistant"]
-        self.offset = 0
+        self.current_prompt = ""
+        self.current_image = None
+        self.current_response = None
         self.skip_next = False
+        self.messages = []  # Add messages list to store conversation history
+
+    def set_prompt(self, prompt, image=None):
+        self.current_prompt = prompt
+        self.current_image = image
+        self.current_response = None
+        # Update messages when setting a new prompt
+        self.messages = [[prompt, None]]
 
-    def append_message(self, role, message):
-        self.messages.append([role, message])
+    def set_response(self, response):
+        self.current_response = response
+        # Update the last message's response when setting a response
+        if self.messages and len(self.messages) > 0:
+            self.messages[-1][-1] = response
+
+    def get_prompt(self):
+        if isinstance(self.current_prompt, tuple):
+            return self.current_prompt[0]
+        return self.current_prompt
+
+    def get_image(self, return_pil=False):
+        if self.current_image:
+            return [self.current_image]
+        if isinstance(self.current_prompt, tuple) and len(self.current_prompt) > 1:
+            if isinstance(self.current_prompt[1], Image.Image):
+                return [self.current_prompt[1]]
+        return None
 
     def to_gradio_chatbot(self):
+        if not self.messages:
+            return []
+            
         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}")
+        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
-    
-    def render_user_message(self, message):
-        if "<image>" in message:
-            return message.replace("<image>", "")
-        return message
 
     def dict(self):
-        # Create a serializable version of messages
-        serialized_messages = []
-        for role, message in self.messages:
-            if isinstance(message, tuple) and len(message) > 1:
-                # If the message contains an image (tuple format)
-                if isinstance(message[1], Image.Image):
-                    # Just keep the text part and ignore the image
-                    serialized_message = (message[0], "[IMAGE_IGNORED]")
-                else:
-                    # For non-image tuples, keep as is
-                    serialized_message = message
-            else:
-                # For non-tuple messages, keep as is
-                serialized_message = message
-            serialized_messages.append([role, serialized_message])
-
+        # Simplified serialization for logging
+        image_info = "[WITH_IMAGE]" if self.current_image is not None else "[NO_IMAGE]"
         return {
-            "messages": serialized_messages,
-            "roles": self.roles,
-            "offset": self.offset,
-            "skip_next": self.skip_next,
+            "prompt": self.get_prompt(),
+            "image": image_info,
+            "response": self.current_response,
+            "messages": [[m[0], "[RESPONSE]" if m[1] else None] for m in self.messages]
         }
-    
-    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 = SimpleConversation()
+        new_conv.current_prompt = self.current_prompt
+        new_conv.current_image = self.current_image
+        new_conv.current_response = self.current_response
         new_conv.skip_next = self.skip_next
+        new_conv.messages = self.messages.copy() if self.messages else []
         return new_conv
 
-default_conversation = Conversation()
+default_conversation = SimpleConversation()
 
 # Model and processor storage
 tokenizer = None
@@ -131,11 +115,6 @@ model = None
 processor = None
 context_len = 8048
 
-# 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:
@@ -158,7 +137,8 @@ def load_model(model_path):
         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,
+                # torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                torch_dtype="auto",
                 device_map="auto" if torch.cuda.is_available() else None
             )
             processor = AutoProcessor.from_pretrained(model_path)
@@ -168,7 +148,8 @@ def load_model(model_path):
         else:
             model = LlavaOnevisionForConditionalGeneration.from_pretrained(
                 model_path,
-                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                # torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                torch_dtype="auto",
                 device_map="auto" if torch.cuda.is_available() else None,
                 trust_remote_code=True
             )
@@ -185,10 +166,10 @@ def load_model(model_path):
 
 def get_model_list():
     models = [
-        'AIML-TUDA/LlavaGuard-v1.2-0.5B-OV-hf',
-        'AIML-TUDA/LlavaGuard-v1.2-7B-OV-hf',
         '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
 
@@ -204,7 +185,6 @@ def run_inference(prompt, image, temperature=0.2, top_p=0.95, max_tokens=512):
     
     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):
@@ -218,39 +198,18 @@ def run_inference(prompt, image, temperature=0.2, top_p=0.95, max_tokens=512):
                     ]
                 }
             ]
-            
             # Process input
-            text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+            text_prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+            image_inputs, video_inputs = process_vision_info(messages)
             inputs = processor(
-                text=[text],
-                images=[image], 
-                padding=True, 
-                return_tensors="pt"
+                text=[text_prompt],
+                images=image_inputs,
+                videos=video_inputs,
+                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]
+            inputs = inputs.to("cuda")
+
             
         # Otherwise assume it's a LlavaGuard model
         else:
@@ -263,39 +222,37 @@ def run_inference(prompt, image, temperature=0.2, top_p=0.95, max_tokens=512):
                     ],
                 },
             ]
-
             text_prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
+            inputs = processor(text=text_prompt, images=image, return_tensors="pt")
 
+        inputs = {k: v.to('cuda') for k, v in inputs.items()}
 
-            # Process input for LlavaGuard models
-            inputs = processor(text=text_prompt, images=image, return_tensors="pt")
+        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 = generated_ids[0, inputs["input_ids"].shape[1]:]
+        response = processor.decode(
+            generated_ids_trimmed,
+            skip_special_tokens=True,
+            # clean_up_tokenization_spaces=False
+        )
+        print(response)
 
-            
-            # 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: {e}")
-        return f"Error during inference: {e}"
+        import traceback
+        error_msg = f"Error during inference: {str(e)}\n{traceback.format_exc()}"
+        print(error_msg)
+        logger.error(error_msg)
+        return f"Error processing image. Please try again."
 
 # Gradio UI functions
 get_window_url_params = """
@@ -359,10 +316,17 @@ def flag_last_response(state, model_selector, request: gr.Request):
 
 def regenerate(state, image_process_mode, request: gr.Request):
     logger.info(f"regenerate. ip: {request.client.host}")
-    state.messages[-1][-1] = None
-    prev_human_msg = state.messages[-2]
-    if type(prev_human_msg[1]) in (tuple, list):
-        prev_human_msg[1] = (*prev_human_msg[1][:2], image_process_mode)
+    if state.messages and len(state.messages) > 0:
+        state.messages[-1][-1] = None
+        if len(state.messages) > 1:
+            prev_human_msg = state.messages[-2]
+            if isinstance(prev_human_msg[0], tuple) and len(prev_human_msg[0]) >= 2:
+                # Handle image process mode for previous message if it's a tuple with image
+                new_msg = list(prev_human_msg)
+                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
 
@@ -378,15 +342,19 @@ 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:
-        if '<image>' not in text:
-            text = text + '\n<image>'
-        text = (text, image, image_process_mode)
         state = default_conversation.copy()
-    state = clear_conv(state)
-    state.append_message(state.roles[0], text)
-    state.append_message(state.roles[1], None)
+        
+    # 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):
@@ -399,47 +367,50 @@ 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_images(return_pil=True)
+    all_images = state.get_image(return_pil=True)
     
     if not all_images:
-        state.messages[-1][-1] = "Error: No image provided"
+        if not state.messages:
+            state.messages = [["Error: No image provided", None]]
+        else:
+            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_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_val}.jpg")
-        if not os.path.isfile(filename):
-            os.makedirs(os.path.dirname(filename), exist_ok=True)
-            image.save(filename)
-    
     # 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
+    
+    # Update the response in the conversation state
+    if not state.messages:
+        state.messages = [[prompt, output]]
+    else:
+        state.messages[-1][-1] = output
+    state.current_response = output
 
     yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 5
 
     finish_tstamp = time.time()
     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_selector,
-            "start": round(start_tstamp, 4),
-            "finish": round(finish_tstamp, 4),
-            "state": state.dict(),
-            "images": all_image_hash,
-            "ip": request.client.host,
-        }
-        fout.write(json.dumps(data) + "\n")
+    try:
+        with open(get_conv_log_filename(), "a") as fout:
+            data = {
+                "tstamp": round(finish_tstamp, 4),
+                "type": "chat",
+                "model": model_selector,
+                "start": round(start_tstamp, 4),
+                "finish": round(finish_tstamp, 4),
+                "state": state.dict(),
+                "images": ['image'],
+                "ip": request.client.host,
+            }
+            fout.write(json.dumps(data) + "\n")
+    except Exception as e:
+        logger.error(f"Error writing log: {str(e)}")
 
 # UI Components
 title_markdown = """
@@ -666,8 +637,9 @@ if __name__ == "__main__":
         ).launch(
             server_name=args.host,
             server_port=args.port,
-            share=True
+            share=args.share
         )
     except Exception as e:
         logger.error(f"Error launching demo: {e}")
         sys.exit(1)
+