Update app.py

app.py

@@ -1,269 +1,192 @@
+import httpx
+import asyncio
 import os
-import numpy as np
-from PIL import Image
-import requests
-import time
-import multiprocessing
+import uuid
+from pathlib import Path
+from typing import Optional, Dict, Any
 import json
-import sys
-from typing import Tuple, List, Dict, Any
+from datetime import datetime
 
-# Add Florence model path to Python path
-florence_path = os.path.join(os.path.dirname(__file__), 'florence-2-large')
-sys.path.append(florence_path)
-
-try:
-    from processing_florence2 import Florence2Processor
-    from configuration_florence2 import Florence2Config
-    import torch
-    import torch.nn.functional as F
-    
-    # Initialize processor with local files
-    config = Florence2Config.from_json_file(os.path.join(florence_path, 'config.json'))
-    processor = Florence2Processor(config)
-    HAVE_PROCESSOR = True
-    print("Successfully loaded Florence processor")
-except Exception as e:
-    print(f"Warning: Could not load Florence processor: {e}")
-    print("Using basic output interpretation")
-    HAVE_PROCESSOR = False
+class MiddlewareClient:
+    def __init__(self, base_url: str = "https://fred808-vssee.hf.space"):
+        self.base_url = base_url.rstrip('/')
+        self.client = httpx.AsyncClient(timeout=30.0)  # 30 second timeout
+        self.requester_id = str(uuid.uuid4())  # Unique ID for this client
+        self.download_dir = Path("downloads")
+        self.download_dir.mkdir(exist_ok=True)
+        
+        # Keep track of our current locks
+        self.current_course: Optional[str] = None
+        self.current_image: Optional[str] = None
+        
+        # Statistics
+        self.stats = {
+            "downloads_started": 0,
+            "downloads_completed": 0,
+            "bytes_downloaded": 0,
+            "start_time": datetime.now().isoformat()
+        }
 
-# Task-specific configuration
-TASK = "<MORE_DETAILED_CAPTION>"  # For detailed image captioning
+    async def close(self):
+        """Close the HTTP client"""
+        await self.client.aclose()
 
-# Model configuration
-MODEL_ID = "microsoft/florence-2-base"
+    async def get_next_course(self) -> Optional[Dict[str, Any]]:
+        """Get next available course"""
+        try:
+            response = await self.client.get(
+                f"{self.base_url}/middleware/next/course",
+                params={"requester_id": self.requester_id}
+            )
+            response.raise_for_status()
+            course_data = response.json()
+            self.current_course = course_data["course_id"]
+            return course_data
+        except httpx.HTTPError as e:
+            if e.response.status_code == 404:
+                print("No more courses available")
+                return None
+            raise
 
+    async def get_next_image(self, course_id: str) -> Optional[Dict[str, Any]]:
+        """Get next available image from a course"""
+        try:
+            response = await self.client.get(
+                f"{self.base_url}/middleware/next/image/{course_id}",
+                params={"requester_id": self.requester_id}
+            )
+            response.raise_for_status()
+            image_data = response.json()
+            self.current_image = image_data["file_id"]
+            return image_data
+        except httpx.HTTPError as e:
+            if e.response.status_code == 404:
+                print(f"No more images available in course {course_id}")
+                return None
+            raise
 
+    async def release_course(self, course_id: str):
+        """Release lock on a course"""
+        try:
+            response = await self.client.post(
+                f"{self.base_url}/middleware/release/course/{course_id}",
+                params={"requester_id": self.requester_id}
+            )
+            response.raise_for_status()
+            self.current_course = None
+        except httpx.HTTPError as e:
+            print(f"Error releasing course {course_id}: {e}")
 
-def load_and_preprocess_image(image_path):
-    # Load image and resize to 32x32
-    img = Image.open(image_path)
-    img = img.resize((32, 32))
-    
-    # Convert to numpy array and normalize to [0,1]
-    img_array = np.array(img).astype(np.float32) / 255.0
-    
-    # Ensure array has shape (32, 32, 3)
-    if len(img_array.shape) == 2:
-        img_array = np.stack([img_array] * 3, axis=-1)
-    
-    # Add batch dimension
-    img_array = img_array[np.newaxis, ...]
-    
-    # Convert tensor to list of single-element lists for API
-    tensor_data = [[float(x)] for x in img_array.flatten()]
-    
-    return tensor_data
+    async def release_image(self, course_id: str, file_id: str):
+        """Release lock on an image"""
+        try:
+            response = await self.client.post(
+                f"{self.base_url}/middleware/release/image/{course_id}/{file_id}",
+                params={"requester_id": self.requester_id}
+            )
+            response.raise_for_status()
+            self.current_image = None
+        except httpx.HTTPError as e:
+            print(f"Error releasing image {file_id}: {e}")
 
-def run_inference(args: Tuple[str, str, int]) -> dict:
-    """Run inference on a specific server with given chunk ID."""
-    server_url, image_path, chunk_id = args
-    
-    try:
-        print(f"\nProcessing server {server_url} with chunk {chunk_id}...")
+    async def download_file(self, course: str, file_id: str) -> bool:
+        """Download a file to local storage"""
+        save_path = self.download_dir / course / file_id
+        save_path.parent.mkdir(exist_ok=True)
         
-        # Load and preprocess image
-        input_tensor = load_and_preprocess_image(image_path)
-        
-        # Prepare request data
-        data = {
-            "inputs": input_tensor
-        }
-        
-        # Send request with timeout
-        print(f"Sending request to {server_url}/compute/{chunk_id}")
-        start_time = time.time()
-        response = requests.post(
-            f"{server_url}/compute/{chunk_id}",
-            json=data,
-            headers={"Content-Type": "application/json"},
-            timeout=10
-        )
-        
-        inference_time = time.time() - start_time
-        
-        if response.status_code == 200:
-            result = response.json()
-            return {
-                "server": server_url,
-                "chunk_id": chunk_id,
-                "success": True,
-                "time": inference_time,
-                "result": result
-            }
-        else:
-            error_msg = f"HTTP {response.status_code}"
-            if hasattr(response, 'text'):
-                error_msg += f": {response.text}"
-            return {
-                "server": server_url,
-                "chunk_id": chunk_id,
-                "success": False,
-                "error": error_msg,
-                "time": inference_time
-            }
-            
-    except Exception as e:
-        return {
-            "server": server_url,
-            "chunk_id": chunk_id,
-            "success": False,
-            "error": str(e),
-            "time": time.time() - start_time if 'start_time' in locals() else None
-        }
-
-def process_model_outputs(outputs, original_shape=(1, -1, 51289)):
-    """Process model outputs using Florence processor for sequence generation."""
-    # Convert outputs to numpy array
-    outputs_array = np.array([x[0] for x in outputs])
-    
-    if HAVE_PROCESSOR:
         try:
-            # Reshape logits to [batch, seq_len, vocab_size]
-            logits = outputs_array.reshape(original_shape)
+            response = await self.client.get(
+                f"{self.base_url}/download",
+                params={"course": course, "file": file_id}
+            )
+            response.raise_for_status()
             
-            if torch.is_tensor(logits):
-                # Use torch operations if available
-                token_ids = torch.argmax(logits, dim=-1)
-            else:
-                # Fallback to numpy
-                token_ids = np.argmax(logits, axis=-1)
+            self.stats["downloads_started"] += 1
             
-            # Decode tokens to text
-            text = processor.batch_decode(token_ids, skip_special_tokens=True)
+            async with await aiofiles.open(save_path, 'wb') as f:
+                async for chunk in response.aiter_bytes():
+                    await f.write(chunk)
+                    self.stats["bytes_downloaded"] += len(chunk)
             
-            # Post-process for the specific task
-            processed_text = processor.post_process_generation(
-                text[0] if isinstance(text, list) else text,
-                task=TASK
-            )
+            self.stats["downloads_completed"] += 1
+            return True
             
-            return {
-                'text': processed_text,
-                'tokens': token_ids.tolist() if torch.is_tensor(token_ids) else token_ids.tolist(),
-                'logits_shape': logits.shape,
-                'distribution': {
-                    'min': float(outputs_array.min()),
-                    'max': float(outputs_array.max()),
-                    'mean': float(outputs_array.mean()),
-                    'std': float(outputs_array.std())
-                }
-            }
         except Exception as e:
-            print(f"Warning: Error in sequence processing: {e}")
-    
-    # Fallback to basic statistics if processor not available
-    return {
-        'overall_mean': float(outputs_array.mean()),
-        'overall_std': float(outputs_array.std()),
-        'shape': outputs_array.shape,
-        'distribution': {
-            'min': float(outputs_array.min()),
-            'max': float(outputs_array.max()),
-            'median': float(np.median(outputs_array))
-        }
-    }
+            print(f"Error downloading {file_id}: {e}")
+            return False
 
-def process_results(results):
-    """Process and combine results from all servers."""
-    # Filter successful results
-    successful_results = [r for r in results if r['success']]
-    if not successful_results:
-        print("\nError: No servers returned successful results")
-        return
-    
-    # Sort successful results by chunk ID
-    successful_results.sort(key=lambda x: x['chunk_id'])
-    
-    print(f"\nModel Output Analysis ({len(successful_results)}/{len(results)} servers succeeded):")
-    print("-" * 80)
-    
-    # Get total sequence length from all chunks
-    total_outputs = []
-    for result in successful_results:
-        total_outputs.extend(result['result']['outputs'])
-    
-    # Process the combined sequence
-    print("\nProcessing complete sequence...")
-    analysis = process_model_outputs(total_outputs, original_shape=(1, -1, 51289))
-    
-    if 'text' in analysis:
-        print("\nGenerated Description:")
-        print("-" * 80)
-        print(analysis['text'])
-        
-        print("\nSequence Statistics:")
-        print(f"- Logits shape: {analysis['logits_shape']}")
-        print(f"- Distribution:")
-        for key, value in analysis['distribution'].items():
-            print(f"  {key}: {value:.4f}")
-    else:
-        print("\nBasic Analysis (Florence processor not available):")
-        print(f"- Sequence length: {len(total_outputs)}")
-        print(f"- Overall activation: {analysis['overall_mean']:.4f} ± {analysis['overall_std']:.4f}")
-        print("\nValue Distribution:")
-        for key, value in analysis['distribution'].items():
-            print(f"- {key}: {value:.4f}")
-    
-    # Check server consistency
-    if len(successful_results) > 1:
-        all_outputs = [np.array([x[0] for x in r['result']['outputs']]) 
-                      for r in successful_results]
-        differences = [np.max(np.abs(all_outputs[0] - tensor)) 
-                      for tensor in all_outputs[1:]]
-        
-        print("\nServer Consistency:")
-        if np.max(differences) < 1e-6:
-            print("Successful servers provided identical results")
-        else:
-            print(f"Variations detected between servers (max diff: {np.max(differences):.6f})")
-    
-    # Print timing summary
-    successful_times = [r['time'] for r in successful_results]
-    print(f"\nProcessing Time Summary:")
-    print(f"- Average: {np.mean(successful_times):.2f}s")
-    print(f"- Range: {min(successful_times):.2f}s - {max(successful_times):.2f}s")
+    def save_stats(self):
+        """Save download statistics"""
+        self.stats["end_time"] = datetime.now().isoformat()
+        stats_file = self.download_dir / "download_stats.json"
+        with open(stats_file, 'w') as f:
+            json.dump(self.stats, f, indent=2)
 
-def main():
-    # Server configurations with their respective chunk IDs
-    servers = [
-        ("https://fred808-ilob.hf.space", 0),
-        ("https://fred808-tserv.hf.space", 1),
-        ("https://fred808-tserve2.hf.space", 2)
-    ]
+    async def download_all(self, max_courses: int = None, max_files: int = None):
+        """Download all available files with optional limits"""
+        try:
+            courses_processed = 0
+            files_downloaded = 0
+            
+            while True:
+                if max_courses and courses_processed >= max_courses:
+                    print(f"Reached maximum courses limit ({max_courses})")
+                    break
+                    
+                course_data = await self.get_next_course()
+                if not course_data:
+                    print("No more courses available")
+                    break
+                
+                course_id = course_data["course_id"]
+                print(f"\nProcessing course: {course_id}")
+                courses_processed += 1
+                
+                course_files = 0
+                while True:
+                    if max_files and files_downloaded >= max_files:
+                        print(f"Reached maximum files limit ({max_files})")
+                        break
+                        
+                    image_data = await self.get_next_image(course_id)
+                    if not image_data:
+                        break
+                    
+                    file_id = image_data["file_id"]
+                    print(f"Downloading: {file_id}")
+                    
+                    if await self.download_file(course_id, file_id):
+                        files_downloaded += 1
+                        course_files += 1
+                        print(f"Successfully downloaded: {file_id}")
+                    
+                    await self.release_image(course_id, file_id)
+                
+                print(f"Completed course {course_id} - Downloaded {course_files} files")
+                await self.release_course(course_id)
+            
+            print(f"\nDownload complete!")
+            print(f"Processed {courses_processed} courses")
+            print(f"Downloaded {files_downloaded} files")
+            print(f"Total bytes: {self.stats['bytes_downloaded']:,}")
+            
+        finally:
+            self.save_stats()
+            await self.close()
+
+async def main():
+    # Create downloads directory
+    Path("downloads").mkdir(exist_ok=True)
     
-    # Image path - using the same image for all servers
-    image_path = "sample_task/test1.png"
-    print(f"\nTesting with image: {image_path}")
+    client = MiddlewareClient()
     
-    # Create process pool
-    with multiprocessing.Pool() as pool:
-        # Prepare arguments for each server
-        args = [(server_url, image_path, chunk_id) for server_url, chunk_id in servers]
-        
-        # Run inference in parallel
-        print("\nStarting parallel inference across all servers...")
-        results = pool.map(run_inference, args)
-        
-        # Display individual server results
-        print("\nServer Results:")
-        print("-" * 80)
-        for result in results:
-            print(f"\nServer: {result['server']}")
-            print(f"Chunk ID: {result['chunk_id']}")
-            print(f"Success: {result['success']}")
-            print(f"Time: {result['time']:.4f}s" if result['time'] else "Time: N/A")
-            
-            if result['success']:
-                print(f"Output shape: {len(result['result']['outputs'])} elements")
-                print("First few outputs:", result['result']['outputs'][:5])
-            else:
-                print(f"Error: {result['error']}")
-            print("-" * 80)
-        
-        # Process and display combined results
-        process_results(results)
+    try:
+        # Download 2 courses with up to 5 files each as an example
+        await client.download_all(max_courses=2, max_files=10)
+    except KeyboardInterrupt:
+        print("\nDownload interrupted by user")
+    finally:
+        await client.close()
 
 if __name__ == "__main__":
-    main()
+    asyncio.run(main())