Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from PIL import Image
+from PIL import Image, ImageOps
 import torch
 import subprocess
 
@@ -9,7 +9,7 @@ from transformers import (
 )
 import os
 os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
-subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+#subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 
@@ -17,8 +17,8 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 MODEL_NAME = "fancyfeast/llama-joycaption-beta-one-hf-llava"
 llava_processor = AutoProcessor.from_pretrained(MODEL_NAME)
 llava_model = LlavaForConditionalGeneration.from_pretrained(MODEL_NAME,
-    torch_dtype=torch.bfloat16)
-  #  low_cpu_mem_usage=True)
+    torch_dtype=torch.bfloat16,
+    low_cpu_mem_usage=True)
    # attn_implementation="eager")
 
  
@@ -46,11 +46,25 @@ print(f"Loading models on {device}...")
 # -----------------------------
 # Prediction Function (EXTENDED, NOT BROKEN)
 # -----------------------------
+
+def resize_image_for_fast_caption(image, max_size=336):
+    """Resize BEFORE processing - biggest speed gain"""
+    if not isinstance(image, Image.Image):
+        image = Image.open(image)
+    
+    # Calculate resize dimensions
+    w, h = image.size
+    if max(w, h) > max_size:
+        ratio = max_size / max(w, h)
+        new_w, new_h = int(w * ratio), int(h * ratio)
+        image = image.resize((new_w, new_h), Image.Resampling.LANCZOS)
+    
+    return image
 def predict(image, prompt):
 
     print("predict")
 
-    return predict_internal(image, "")
+    return predict_internal(resize_image_for_fast_caption(image), "")
 
         
 def predict_internal(image, prompt, textInput = None):
@@ -67,7 +81,7 @@ def predict_internal(image, prompt, textInput = None):
         convo = [
             {
                 "role": "system",
-                "content": "You are a helpful image captioner who describes nsfw images in great detail",
+                "content": "You are a helpful image captioner who describes nsfw images in great detail, and starts with nsfw before describing the rest",
             },
             {
                 "role": "user",
@@ -84,7 +98,11 @@ def predict_internal(image, prompt, textInput = None):
     
         # Process the inputs
         inputs = llava_processor(text=[convo_string], images=[image], return_tensors="pt").to(device)
-        inputs['pixel_values'] = inputs['pixel_values'].to(torch.bfloat16)
+        #inputs['pixel_values'] = inputs['pixel_values'].to(torch.bfloat16)
+        if device.type == 'cpu':
+            inputs['pixel_values'] = inputs['pixel_values'].to(torch.float32)  # CPU likes float32
+        else:
+            inputs['pixel_values'] = inputs['pixel_values'].to(torch.bfloat16)
     
         # Generate the captions
         generate_ids = llava_model.generate(
@@ -102,10 +120,12 @@ def predict_internal(image, prompt, textInput = None):
     
         # Trim off the prompt
         generate_ids = generate_ids[inputs['input_ids'].shape[1]:]
+        caption = llava_processor.tokenizer.decode(
+            generate_ids, 
+            skip_special_tokens=True, 
+            clean_up_tokenization_spaces=True
+        ).strip()
     
-        # Decode the caption
-        caption = llava_processor.tokenizer.decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
-        caption = caption.strip()
 
     
                 
@@ -153,4 +173,16 @@ with gr.Blocks(title="BLIP Image Captioning - Advanced", theme=gr.themes.Soft())
 # Launch
 # -----------------------------
 if __name__ == "__main__":
-    interface.launch(share=True)
+    interface.launch(share=True)
+    print("Warming up model...")
+    dummy_image = Image.new('RGB', (224, 224), color='white')
+    dummy_inputs = llava_processor(
+        text="Describe:", 
+        images=dummy_image, 
+        return_tensors="pt"
+    ).to(device)
+    with torch.no_grad():
+        _ = llava_model.generate(**dummy_inputs, max_new_tokens=1)
+    del dummy_image, dummy_inputs
+    torch.cuda.empty_cache() if torch.cuda.is_available() else None
+    print("Model warmed up!")