Update app.py

app.py

@@ -1,10 +1,193 @@
 import gradio as gr
+import torch
+from transformers import AutoProcessor, AutoModelForVision2Seq
+from PIL import Image
+import io
+import json
+import time
+import os
+import hashlib
 
-with gr.Blocks(fill_height=True) as demo:
-    with gr.Sidebar():
-        gr.Markdown("# Inference Provider")
-        gr.Markdown("This Space showcases the meta-llama/Llama-4-Scout-17B-16E-Instruct model, served by the together API. Sign in with your Hugging Face account to use this API.")
-        button = gr.LoginButton("Sign in")
-    gr.load("models/meta-llama/Llama-4-Scout-17B-16E-Instruct", accept_token=button, provider="together")
+# Global variables for model and processor
+model = None
+processor = None
+
+# Initialize model and processor
+def load_model():
+    global model, processor
+    if model is None:
+        try:
+            print("Loading Llama 4 Scout model...")
+            processor = AutoProcessor.from_pretrained("meta-llama/Llama-4-Scout-17B-16E-Instruct")
+            model = AutoModelForVision2Seq.from_pretrained(
+                "meta-llama/Llama-4-Scout-17B-16E-Instruct",
+                torch_dtype=torch.float16,
+                device_map="auto"
+            )
+            print("Model loaded successfully!")
+        except Exception as e:
+            print(f"Error loading model: {e}")
+            raise
+    return model, processor
+
+# Simple caching mechanism
+cache = {}
+
+def compute_image_hash(image):
+    """Compute a hash for an image to use as cache key"""
+    # Resize to small dimensions to ensure hash is based on content, not size
+    image = image.resize((100, 100), Image.LANCZOS)
+    
+    # Convert to bytes
+    img_byte_arr = io.BytesIO()
+    image.save(img_byte_arr, format='PNG')
+    img_byte_arr = img_byte_arr.getvalue()
+    
+    # Compute hash
+    return hashlib.md5(img_byte_arr).hexdigest()
+
+def verify_document(img, doc_type, verification_info):
+    """Verify a document using Llama 4 Scout"""
+    if img is None:
+        return "Please upload an image"
+    
+    # Compute image hash for caching
+    image_hash = compute_image_hash(img)
+    cache_key = f"verify_{image_hash}_{doc_type}"
+    
+    # Check cache
+    if cache_key in cache:
+        return f"[CACHED] {cache[cache_key]}"
+    
+    try:
+        # Load model
+        model, processor = load_model()
+        
+        # Create prompt
+        prompt = f"""This is a {doc_type} document. 
+Verify if it's authentic and extract the following information: {verification_info}
+Provide your analysis in a structured format."""
+        
+        # Process with model
+        inputs = processor(text=prompt, images=img, return_tensors="pt").to(model.device)
+        outputs = model.generate(**inputs, max_new_tokens=500)
+        result = processor.decode(outputs[0], skip_special_tokens=True)
+        
+        # Save to cache
+        cache[cache_key] = result
+        
+        return result
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+def check_workplace(img, industry):
+    """Check workplace compliance using Llama 4 Scout"""
+    if img is None:
+        return "Please upload an image"
+    
+    # Compute image hash for caching
+    image_hash = compute_image_hash(img)
+    cache_key = f"workplace_{image_hash}_{industry}"
+    
+    # Check cache
+    if cache_key in cache:
+        return f"[CACHED] {cache[cache_key]}"
     
+    try:
+        # Load model
+        model, processor = load_model()
+        
+        # Create prompt
+        prompt = f"""This is a workplace in the {industry} industry.
+Identify any safety or compliance issues visible in this image.
+Focus on:
+1. Safety hazards
+2. Required signage
+3. Proper equipment usage
+4. Workspace organization
+5. Compliance with regulations
+
+Format your response as a detailed assessment with:
+- Issues found (if any)
+- Severity level for each issue
+- Recommendations for correction"""
+        
+        # Process with model
+        inputs = processor(text=prompt, images=img, return_tensors="pt").to(model.device)
+        outputs = model.generate(**inputs, max_new_tokens=800)
+        result = processor.decode(outputs[0], skip_special_tokens=True)
+        
+        # Save to cache
+        cache[cache_key] = result
+        
+        return result
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# Create Gradio interface
+with gr.Blocks(title="StaffManager AI Assistant") as demo:
+    gr.Markdown("# StaffManager AI Assistant")
+    gr.Markdown("This Space provides AI capabilities for StaffManager using Llama 4 Scout.")
+    
+    with gr.Tab("Document Verification"):
+        with gr.Row():
+            with gr.Column():
+                doc_image = gr.Image(type="pil", label="Upload Document")
+                doc_type = gr.Dropdown(
+                    ["identification", "tax", "employment", "policy"], 
+                    label="Document Type",
+                    value="identification"
+                )
+                verification_info = gr.Textbox(
+                    label="Verification Data (JSON)",
+                    value='{"name": "John Doe", "id_number": "ABC123456"}'
+                )
+                verify_button = gr.Button("Verify Document")
+            with gr.Column():
+                doc_result = gr.Textbox(label="Verification Result", lines=10)
+        
+        verify_button.click(
+            fn=verify_document,
+            inputs=[doc_image, doc_type, verification_info],
+            outputs=[doc_result]
+        )
+    
+    with gr.Tab("Workplace Compliance"):
+        with gr.Row():
+            with gr.Column():
+                workplace_image = gr.Image(type="pil", label="Upload Workplace Image")
+                industry_type = gr.Dropdown(
+                    ["retail", "restaurant", "healthcare", "manufacturing"], 
+                    label="Industry",
+                    value="retail"
+                )
+                check_button = gr.Button("Check Compliance")
+            with gr.Column():
+                compliance_result = gr.Textbox(label="Compliance Assessment", lines=10)
+        
+        check_button.click(
+            fn=check_workplace,
+            inputs=[workplace_image, industry_type],
+            outputs=[compliance_result]
+        )
+    
+    with gr.Tab("About"):
+        gr.Markdown("""
+        ## About StaffManager AI Assistant
+        
+        This Space uses the Llama 4 Scout model to provide AI capabilities for StaffManager:
+        
+        - **Document Verification**: Verify and extract information from documents
+        - **Workplace Compliance**: Identify safety and compliance issues in workplace images
+        
+        The model is loaded on demand and results are cached for better performance.
+        
+        ### Model Information
+        
+        - Model: meta-llama/Llama-4-Scout-17B-16E-Instruct
+        - Type: Multimodal (image + text)
+        - Size: 17B parameters
+        """)
+
+# Launch the app
 demo.launch()