wip

app.py

@@ -1,18 +1,10 @@
-from fastapi import FastAPI #, HTTPException, status, Security
-# from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
-from decouple import config
+from fastapi import FastAPI
 
 from src.encoder import FashionCLIPEncoder
 from src.models import TextRequest, ImageRequest, Response
 
 
-# security = HTTPBearer()
 encoder = FashionCLIPEncoder()
-
-
-API_TOKEN = config("API_TOKEN")
-
-
 app = FastAPI()
 
 
@@ -30,15 +22,8 @@ async def root():
 
 @app.post("/encode_texts")
 async def encode_texts(
-    request: TextRequest, 
-    # credentials: HTTPAuthorizationCredentials = Security(security)
+    request: TextRequest,
 ) -> Response:
-    # if credentials.credentials != API_TOKEN:
-    #     raise HTTPException(
-    #         status_code=status.HTTP_401_UNAUTHORIZED,
-    #         detail="Invalid authentication token",
-    #     )
-
     embeddings = encoder.encode_text(request.texts)
     response = Response(embeddings=embeddings)
 
@@ -48,14 +33,7 @@ async def encode_texts(
 @app.post("/encode_images")
 async def encode_images(
     request: ImageRequest,
-    # credentials: HTTPAuthorizationCredentials = Security(security),
 ) -> Response:
-    # if credentials.credentials != API_TOKEN:
-    #     raise HTTPException(
-    #         status_code=status.HTTP_401_UNAUTHORIZED,
-    #         detail="Invalid authentication token",
-    #     )
-
     images = request.download()
     embeddings = encoder.encode_images(images)
     response = Response(embeddings=embeddings)

src/encoder.py

@@ -11,12 +11,12 @@ MODEL_NAME = "Marqo/marqo-fashionCLIP"
 class FashionCLIPEncoder:
     def __init__(self):
         self.device = torch.device("cpu")
-        
+
         self.processor = AutoProcessor.from_pretrained(
-            MODEL_NAME, 
+            MODEL_NAME,
             trust_remote_code=True,
         )
-        
+
         self.model = AutoModel.from_pretrained(
             MODEL_NAME,
             trust_remote_code=True,
@@ -51,4 +51,4 @@ class FashionCLIPEncoder:
         return self.model.get_text_features(**batch).detach().cpu().numpy().tolist()
 
     def _encode_images(self, batch: Dict) -> List[List[float]]:
-        return self.model.get_image_features(**batch).detach().cpu().numpy().tolist()
+        return self.model.get_image_features(**batch).detach().cpu().numpy().tolist()

src/utils.py

@@ -26,39 +26,41 @@ def analyze_model_parameters(model: torch.nn.Module) -> Dict:
     total_params = 0
     param_types = set()
     param_type_counts = {}
-    
+
     for param in model.parameters():
         total_params += param.numel()
         dtype = param.dtype
         param_types.add(dtype)
         param_type_counts[dtype] = param_type_counts.get(dtype, 0) + param.numel()
-    
+
     results = {
         "total_params": total_params,
         "param_types": {},
         "device_info": {
             "device": next(model.parameters()).device,
-            "cuda_available": torch.cuda.is_available()
-        }
+            "cuda_available": torch.cuda.is_available(),
+        },
     }
-    
+
     for dtype in param_types:
         count = param_type_counts[dtype]
         percentage = (count / total_params) * 100
         memory_bytes = count * torch.finfo(dtype).bits // 8
         memory_mb = memory_bytes / (1024 * 1024)
-        
+
         results["param_types"][str(dtype)] = {
             "count": count,
             "percentage": percentage,
-            "memory_mb": memory_mb
+            "memory_mb": memory_mb,
         }
-    
+
     if torch.cuda.is_available():
-        results["device_info"].update({
-            "cuda_device": torch.cuda.get_device_name(0),
-            "cuda_memory_allocated_mb": torch.cuda.memory_allocated(0) / 1024**2,
-            "cuda_memory_cached_mb": torch.cuda.memory_reserved(0) / 1024**2
-        })
-    
-    return results
+        results["device_info"].update(
+            {
+                "cuda_device": torch.cuda.get_device_name(0),
+                "cuda_memory_allocated_mb": torch.cuda.memory_allocated(0) / 1024**2,
+                "cuda_memory_cached_mb": torch.cuda.memory_reserved(0) / 1024**2,
+            }
+        )
+
+    return results