app.py

import torch
import os
import glob
from PIL import Image
from transformers import AutoProcessor, AutoModelForImageClassification
import gradio as gr
import pytesseract

def find_model_files():
    """Find model files in the current directory structure"""
    print("=== Searching for model files ===")
    
    # Look for key model files
    config_files = glob.glob("**/config.json", recursive=True)
    model_files = glob.glob("**/pytorch_model.bin", recursive=True) + glob.glob("**/model.safetensors", recursive=True)
    preprocessor_files = glob.glob("**/preprocessor_config.json", recursive=True)
    
    print(f"Found config.json files: {config_files}")
    print(f"Found model weight files: {model_files}")
    print(f"Found preprocessor_config.json files: {preprocessor_files}")
    
    # Find the directory that contains all necessary files
    for config_file in config_files:
        model_dir = os.path.dirname(config_file)
        if not model_dir:  # If config.json is in root
            model_dir = "."
        
        # Check if this directory has all required files
        has_model = any(os.path.dirname(f) == model_dir or (not os.path.dirname(f) and model_dir == ".") for f in model_files)
        has_preprocessor = any(os.path.dirname(f) == model_dir or (not os.path.dirname(f) and model_dir == ".") for f in preprocessor_files)
        
        if has_model and has_preprocessor:
            print(f"Found complete model in directory: {model_dir}")
            return model_dir
        elif has_model:
            print(f"Found model with config but missing preprocessor in: {model_dir}")
            return model_dir  # Try anyway, might work
    
    print("No complete model directory found")
    return None

# Search for model files
MODEL_PATH = find_model_files()
if MODEL_PATH is None:
    MODEL_PATH = "."  # Fallback to current directory
    print("Falling back to current directory")

try:
    # Load model and processor from detected path
    print(f"=== Attempting to load model from: {MODEL_PATH} ===")
    print(f"Current working directory: {os.getcwd()}")
    
    # List all files in the detected model directory
    if MODEL_PATH == ".":
        print("Files in root directory:")
        for item in os.listdir("."):
            if os.path.isfile(item):
                print(f"  File: {item}")
            else:
                print(f"  Directory: {item}/")
                try:
                    sub_files = os.listdir(item)[:5]  # Show first 5 files
                    print(f"    Contains: {sub_files}{'...' if len(os.listdir(item)) > 5 else ''}")
                except:
                    pass
    else:
        print(f"Files in {MODEL_PATH}:")
        print(f"  {os.listdir(MODEL_PATH)}")
    
    # Try to load the model
    print("Loading model...")
    model = AutoModelForImageClassification.from_pretrained(MODEL_PATH, local_files_only=True)
    print("Model loaded successfully!")
    
    print("Loading processor...")
    try:
        processor = AutoProcessor.from_pretrained(MODEL_PATH, local_files_only=True)
        print("Processor loaded successfully!")
    except Exception as proc_error:
        print(f"Error loading processor from local files: {proc_error}")
        print("Attempting to load just the image processor...")
        
        # Try to load just the image processor from your model
        try:
            from transformers import SiglipImageProcessor
            processor = SiglipImageProcessor.from_pretrained(MODEL_PATH, local_files_only=True)
            print("Image processor loaded successfully from local files!")
        except Exception as img_proc_error:
            print(f"Error loading local image processor: {img_proc_error}")
            print("Loading image processor from base SigLIP model...")
            
            # Try to load processor from the base SigLIP model
            try:
                from transformers import SiglipImageProcessor
                processor = SiglipImageProcessor.from_pretrained("google/siglip-base-patch16-224")
                print("Image processor loaded from base SigLIP model!")
            except Exception as base_error:
                print(f"Error loading base processor: {base_error}")
                print("Using CLIP processor as fallback...")
                
                # As a last resort, try to create a minimal processor
                from transformers import CLIPImageProcessor
                processor = CLIPImageProcessor.from_pretrained("openai/clip-vit-base-patch32")
                print("Using CLIP processor as fallback!")
    
    # Get labels - handle case where id2label might not exist
    if hasattr(model.config, 'id2label') and model.config.id2label:
        labels = model.config.id2label
    else:
        # Create generic labels if none exist
        num_labels = model.config.num_labels if hasattr(model.config, 'num_labels') else 1000
        labels = {i: f"class_{i}" for i in range(num_labels)}
    
    print(f"Model loaded successfully. Number of classes: {len(labels)}")
    
except Exception as e:
    print(f"=== ERROR loading model from {MODEL_PATH} ===")
    print(f"Error: {e}")
    print("\n=== Debugging Information ===")
    print("All files in Space:")
    
    def list_all_files(directory=".", prefix=""):
        """Recursively list all files"""
        try:
            items = sorted(os.listdir(directory))
            for item in items:
                item_path = os.path.join(directory, item)
                if os.path.isfile(item_path):
                    size = os.path.getsize(item_path)
                    print(f"{prefix}📄 {item} ({size} bytes)")
                elif os.path.isdir(item_path) and not item.startswith('.'):
                    print(f"{prefix}📁 {item}/")
                    if len(prefix) < 6:  # Limit recursion depth
                        list_all_files(item_path, prefix + "  ")
        except PermissionError:
            print(f"{prefix}❌ Permission denied")
        except Exception as ex:
            print(f"{prefix}❌ Error: {ex}")
    
    list_all_files()
    
    print("\n=== Required Files for Model ===")
    print("✅ config.json - Model configuration")
    print("✅ pytorch_model.bin OR model.safetensors - Model weights") 
    print("✅ preprocessor_config.json - Image processor config")
    print("✅ tokenizer.json (if applicable) - Tokenizer")
    
    print("\n=== Solutions ===")
    print("1. Make sure all model files are uploaded to your Space")
    print("2. Check that files aren't corrupted during upload")
    print("3. Try uploading to a 'model' subfolder")
    print("4. Verify the model was saved correctly during training")
    
    raise

# Classify meme and extract text
def classify_meme(image: Image.Image):
    """
    Classify meme and extract text using OCR
    """
    try:
        # OCR: extract text from image
        extracted_text = pytesseract.image_to_string(image)
        
        # Process image with the model
        inputs = processor(images=image, return_tensors="pt")
        
        # Move inputs to same device as model if needed
        if torch.cuda.is_available() and next(model.parameters()).is_cuda:
            inputs = {k: v.to('cuda') for k, v in inputs.items()}
        
        with torch.no_grad():
            outputs = model(**inputs)
            probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
        
        # Get top predictions
        top_k = min(10, len(labels))  # Show top 10 or all if fewer
        top_probs, top_indices = torch.topk(probs[0], top_k)
        
        predictions = {}
        for i, (prob, idx) in enumerate(zip(top_probs, top_indices)):
            label = labels.get(idx.item(), f"class_{idx.item()}")
            predictions[label] = float(prob)
        
        # Debug prints (these will show in the console/logs)
        print("Extracted Text:", extracted_text.strip())
        print("Top Predictions:", predictions)
        
        return predictions, extracted_text.strip()
        
    except Exception as e:
        print(f"Error in classification: {e}")
        return {"Error": 1.0}, f"Error processing image: {str(e)}"

# Gradio interface
demo = gr.Interface(
    fn=classify_meme,
    inputs=gr.Image(type="pil", label="Upload Meme Image"),
    outputs=[
        gr.Label(num_top_classes=5, label="Meme Classification"),
        gr.Textbox(label="Extracted Text from OCR", lines=3)
    ],
    title="Meme Classifier with OCR",
    description="""
    Upload a meme image to:
    1. Classify its content using your trained SigLIP2_77 model
    2. Extract text using OCR (Optical Character Recognition)
    
    Note: Make sure all model files are properly uploaded to your Space.
    """,
    examples=None,
    allow_flagging="never"
)

if __name__ == "__main__":
    print("Starting Gradio interface...")
    demo.launch(
        server_name="0.0.0.0",  # Allow external connections in HF Spaces
        server_port=7860,       # Standard port for HF Spaces
        share=False             # HF Spaces handles sharing
    )