Update app.py

app.py

@@ -1,184 +1,273 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor
+from transformers import AutoModelForCausalLM, AutoTokenizer, CLIPProcessor, CLIPModel
 from PIL import Image
-import os
+import logging
 import spaces
+import numpy
 
-# Initial setup without loading model to device
-print("Setting up the application...")
+# Setup logging
+logging.basicConfig(level=logging.INFO)
 
-# We'll load the model in the GPU functions to avoid CPU memory issues
-model = None
-tokenizer = AutoTokenizer.from_pretrained("sagar007/Lava_phi")
-processor = AutoProcessor.from_pretrained("openai/clip-vit-base-patch32")
-
-print("Tokenizer and processor loaded successfully!")
-
-# For text-only generation with GPU on demand
-@spaces.GPU
-def generate_text(prompt, max_length=128):
-    try:
-        global model
+class LLaVAPhiModel:
+    def __init__(self, model_id="sagar007/Lava_phi"):
+        self.device = "cuda"
+        self.model_id = model_id
+        logging.info("Initializing LLaVA-Phi model...")
         
-        # Load model if not already loaded
-        if model is None:
-            print("Loading model on first request...")
-            model = AutoModelForCausalLM.from_pretrained(
-                "sagar007/Lava_phi",
-                torch_dtype=torch.float16,  # Use float16 on GPU
-                device_map="auto"  # This will put the model on GPU automatically
-            )
-            print("Model loaded successfully!")
+        # Initialize tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
         
-        inputs = tokenizer(f"human: {prompt}\ngpt:", return_tensors="pt").to(model.device)
+        try:
+            # Use CLIPProcessor directly instead of AutoProcessor
+            self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+            logging.info("Successfully loaded CLIP processor")
+        except Exception as e:
+            logging.error(f"Failed to load CLIP processor: {str(e)}")
+            self.processor = None
         
-        # Generate with GPU
-        with torch.no_grad():
-            outputs = model.generate(
-                **inputs, 
-                max_new_tokens=max_length,
-                do_sample=True,
-                temperature=0.7,
-                top_p=0.9,
+        self.history = []
+        self.model = None
+        self.clip = None
+
+    @spaces.GPU
+    def ensure_models_loaded(self):
+        """Ensure models are loaded in GPU context"""
+        if self.model is None:
+            # Load main model with updated quantization config
+            from transformers import BitsAndBytesConfig
+            quantization_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype=torch.float16,
+                bnb_4bit_use_double_quant=True,
+                bnb_4bit_quant_type="nf4"
             )
-        
-        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        
-        # Extract only the model's response
-        if "gpt:" in generated_text:
-            generated_text = generated_text.split("gpt:", 1)[1].strip()
             
-        return generated_text
-    except Exception as e:
-        # Capture and return any errors
-        return f"Error generating text: {str(e)}"
+            try:
+                self.model = AutoModelForCausalLM.from_pretrained(
+                    self.model_id,
+                    quantization_config=quantization_config,
+                    device_map="auto",
+                    torch_dtype=torch.bfloat16,
+                    trust_remote_code=True
+                )
+                self.model.config.pad_token_id = self.tokenizer.eos_token_id
+                logging.info("Successfully loaded main model")
+            except Exception as e:
+                logging.error(f"Failed to load main model: {str(e)}")
+                raise
+
+        if self.clip is None:
+            try:
+                # Use CLIPModel directly instead of AutoModel
+                self.clip = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(self.device)
+                logging.info("Successfully loaded CLIP model")
+            except Exception as e:
+                logging.error(f"Failed to load CLIP model: {str(e)}")
+                self.clip = None
+
+    @spaces.GPU
+    def process_image(self, image):
+        """Process image through CLIP if available"""
+        try:
+            self.ensure_models_loaded()
+            
+            if self.clip is None or self.processor is None:
+                logging.warning("CLIP model or processor not available")
+                return None
+            
+            # Convert image to correct format
+            if isinstance(image, str):
+                image = Image.open(image)
+            elif isinstance(image, numpy.ndarray):
+                image = Image.fromarray(image)
+            
+            # Ensure image is in RGB mode
+            if image.mode != 'RGB':
+                image = image.convert('RGB')
+            
+            with torch.no_grad():
+                try:
+                    # Process image with error handling
+                    image_inputs = self.processor(images=image, return_tensors="pt")
+                    image_features = self.clip.get_image_features(
+                        pixel_values=image_inputs.pixel_values.to(self.device)
+                    )
+                    logging.info("Successfully processed image through CLIP")
+                    return image_features
+                except Exception as e:
+                    logging.error(f"Error during image processing: {str(e)}")
+                    return None
+        except Exception as e:
+            logging.error(f"Error in process_image: {str(e)}")
+            return None
+
+    @spaces.GPU(duration=120)
+    def generate_response(self, message, image=None):
+        try:
+            self.ensure_models_loaded()
+            
+            if image is not None:
+                image_features = self.process_image(image)
+                has_image = image_features is not None
+                if not has_image:
+                    message = "Note: Image processing is not available - continuing with text only.\n" + message
+                
+                prompt = f"human: {'<image>' if has_image else ''}\n{message}\ngpt:"
+                context = ""
+                for turn in self.history[-3:]:
+                    context += f"human: {turn[0]}\ngpt: {turn[1]}\n"
+                
+                full_prompt = context + prompt
+                inputs = self.tokenizer(
+                    full_prompt, 
+                    return_tensors="pt", 
+                    padding=True,
+                    truncation=True,
+                    max_length=512
+                )
+                inputs = {k: v.to(self.device) for k, v in inputs.items()}
+                
+                if has_image:
+                    inputs["image_features"] = image_features
+                
+                with torch.no_grad():
+                    outputs = self.model.generate(
+                        **inputs,
+                        max_new_tokens=256,
+                        min_length=20,
+                        temperature=0.7,
+                        do_sample=True,
+                        top_p=0.9,
+                        top_k=40,
+                        repetition_penalty=1.5,
+                        no_repeat_ngram_size=3,
+                        use_cache=True,
+                        pad_token_id=self.tokenizer.pad_token_id,
+                        eos_token_id=self.tokenizer.eos_token_id
+                    )
+            else:
+                prompt = f"human: {message}\ngpt:"
+                context = ""
+                for turn in self.history[-3:]:
+                    context += f"human: {turn[0]}\ngpt: {turn[1]}\n"
+                
+                full_prompt = context + prompt
+                inputs = self.tokenizer(
+                    full_prompt,
+                    return_tensors="pt",
+                    padding=True,
+                    truncation=True,
+                    max_length=512
+                )
+                inputs = {k: v.to(self.device) for k, v in inputs.items()}
+                
+                with torch.no_grad():
+                    outputs = self.model.generate(
+                        **inputs,
+                        max_new_tokens=150,
+                        min_length=20,
+                        temperature=0.6,
+                        do_sample=True,
+                        top_p=0.85,
+                        top_k=30,
+                        repetition_penalty=1.8,
+                        no_repeat_ngram_size=4,
+                        use_cache=True,
+                        pad_token_id=self.tokenizer.pad_token_id,
+                        eos_token_id=self.tokenizer.eos_token_id
+                    )
+            
+            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            
+            # Clean up response
+            if "gpt:" in response:
+                response = response.split("gpt:")[-1].strip()
+            if "human:" in response:
+                response = response.split("human:")[0].strip()
+            if "<image>" in response:
+                response = response.replace("<image>", "").strip()
+            
+            self.history.append((message, response))
+            return response
+            
+        except Exception as e:
+            logging.error(f"Error generating response: {str(e)}")
+            logging.error(f"Full traceback:", exc_info=True)
+            return f"Error: {str(e)}"
+
+    def clear_history(self):
+        self.history = []
+        return None
 
-# For image and text processing with GPU on demand
-@spaces.GPU
-def process_image_and_prompt(image, prompt, max_length=128):
+def create_demo():
     try:
-        if image is None:
-            return "No image provided. Please upload an image."
+        model = LLaVAPhiModel()
         
-        global model
-        
-        # Load model if not already loaded
-        if model is None:
-            print("Loading model on first request...")
-            model = AutoModelForCausalLM.from_pretrained(
-                "sagar007/Lava_phi",
-                torch_dtype=torch.float16,  # Use float16 on GPU
-                device_map="auto"  # This will put the model on GPU automatically
+        with gr.Blocks(css="footer {visibility: hidden}") as demo:
+            gr.Markdown(
+                """
+                # LLaVA-Phi Demo (ZeroGPU)
+                Chat with a vision-language model that can understand both text and images.
+                """
             )
-            print("Model loaded successfully!")
-        
-        # Process image
-        image_tensor = processor(images=image, return_tensors="pt").pixel_values.to(model.device)
-        
-        # Tokenize input with image token
-        inputs = tokenizer(f"human: <image>\n{prompt}\ngpt:", return_tensors="pt").to(model.device)
-        
-        # Generate with GPU
-        with torch.no_grad():
-            outputs = model.generate(
-                input_ids=inputs["input_ids"],
-                attention_mask=inputs["attention_mask"],
-                images=image_tensor,
-                max_new_tokens=max_length,
-                do_sample=True,
-                temperature=0.7,
-                top_p=0.9,
+            
+            chatbot = gr.Chatbot(height=400)
+            with gr.Row():
+                with gr.Column(scale=0.7):
+                    msg = gr.Textbox(
+                        show_label=False,
+                        placeholder="Enter text and/or upload an image",
+                        container=False
+                    )
+                with gr.Column(scale=0.15, min_width=0):
+                    clear = gr.Button("Clear")
+                with gr.Column(scale=0.15, min_width=0):
+                    submit = gr.Button("Submit", variant="primary")
+            
+            image = gr.Image(type="pil", label="Upload Image (Optional)")
+            
+            def respond(message, chat_history, image):
+                if not message and image is None:
+                    return chat_history
+                
+                response = model.generate_response(message, image)
+                chat_history.append((message, response))
+                return "", chat_history
+            
+            def clear_chat():
+                model.clear_history()
+                return None, None
+            
+            submit.click(
+                respond,
+                [msg, chatbot, image],
+                [msg, chatbot],
+            )
+            
+            clear.click(
+                clear_chat,
+                None,
+                [chatbot, image],
+            )
+            
+            msg.submit(
+                respond,
+                [msg, chatbot, image],
+                [msg, chatbot],
             )
-        
-        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        
-        # Extract only the model's response
-        if "gpt:" in generated_text:
-            generated_text = generated_text.split("gpt:", 1)[1].strip()
             
-        return generated_text
+        return demo
     except Exception as e:
-        # Capture and return any errors
-        return f"Error processing image: {str(e)}"
-
-# Create Gradio Interface
-with gr.Blocks(title="LLaVA-Phi: Vision-Language Model") as demo:
-    gr.Markdown("# LLaVA-Phi: Vision-Language Model")
-    gr.Markdown("This model uses ZeroGPU technology - GPU resources are allocated only when generating responses and released afterward.")
-    
-    with gr.Tab("Text Generation"):
-        with gr.Row():
-            with gr.Column():
-                text_input = gr.Textbox(label="Enter your prompt", lines=3, placeholder="What is artificial intelligence?")
-                text_max_length = gr.Slider(minimum=16, maximum=512, value=128, step=8, label="Maximum response length")
-                text_button = gr.Button("Generate")
-            
-            with gr.Column():
-                text_output = gr.Textbox(label="Generated response", lines=8)
-                text_status = gr.Markdown("*Status: Ready*")
-        
-        def text_fn(prompt, max_length):
-            text_status.update("*Status: Generating response...*")
-            try:
-                response = generate_text(prompt, max_length)
-                text_status.update("*Status: Complete*")
-                return response
-            except Exception as e:
-                text_status.update("*Status: Error*")
-                return f"Error: {str(e)}"
-        
-        text_button.click(
-            fn=text_fn,
-            inputs=[text_input, text_max_length],
-            outputs=text_output
-        )
-    
-    with gr.Tab("Image + Text Analysis"):
-        with gr.Row():
-            with gr.Column():
-                image_input = gr.Image(type="pil", label="Upload an image")
-                image_text_input = gr.Textbox(label="Enter your prompt about the image", 
-                                              lines=2, 
-                                              placeholder="Describe this image in detail.")
-                image_max_length = gr.Slider(minimum=16, maximum=512, value=128, step=8, label="Maximum response length")
-                image_button = gr.Button("Analyze")
-            
-            with gr.Column():
-                image_output = gr.Textbox(label="Model response", lines=8)
-                image_status = gr.Markdown("*Status: Ready*")
-        
-        def image_fn(image, prompt, max_length):
-            image_status.update("*Status: Analyzing image...*")
-            try:
-                response = process_image_and_prompt(image, prompt, max_length)
-                image_status.update("*Status: Complete*")
-                return response
-            except Exception as e:
-                image_status.update("*Status: Error*")
-                return f"Error: {str(e)}"
-        
-        image_button.click(
-            fn=image_fn,
-            inputs=[image_input, image_text_input, image_max_length],
-            outputs=image_output
-        )
-    
-    # Example inputs for each tab
-    gr.Examples(
-        examples=["What is the advantage of vision-language models?", 
-                  "Explain how multimodal AI models work.",
-                  "Tell me a short story about robots."],
-        inputs=text_input
-    )
-    
-    # Status indicator
-    with gr.Row():
-        gr.Markdown("*Note: When you click Generate or Analyze, a GPU will be temporarily allocated to process your request and then released. The first request may take longer as the model needs to be loaded.*")
+        logging.error(f"Error creating demo: {str(e)}")
+        raise
 
-# Launch the app
 if __name__ == "__main__":
+    demo = create_demo()
     demo.launch(
-        enable_queue=True,
-        show_error=True
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=True
     )