Update app.py

app.py

@@ -7,7 +7,6 @@ import json
 import time
 import os
 import hashlib
-import huggingface_hub
 from huggingface_hub import login
 
 # Print token information (first few characters only for security)
@@ -24,9 +23,6 @@ try:
 except Exception as e:
     print(f"Error logging in: {e}")
 
-# Set token for all huggingface_hub operations
-huggingface_hub.set_token(token)
-
 # Simple test to verify token works
 try:
     from huggingface_hub import whoami
@@ -74,4 +70,164 @@ def load_model():
             raise
     return model, processor
 
-# Rest of the code remains the same...
+# 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()