Update app.py

app.py

@@ -31,13 +31,42 @@ class ModelCache:
         self.cache_dir = cache_dir
         self.cache_dir.mkdir(parents=True, exist_ok=True)
 
-    def load_model(self, model_id: str, load_func: callable, cache_name: str) -> Any:
-        try:
-            logger.info(f"Loading {cache_name}")
-            return load_func(model_id)
-        except Exception as e:
-            logger.error(f"Error loading model {cache_name}: {str(e)}")
-            raise
+    import requests
+
+def _load_banglaclip_model(self, weights_path: str) -> CLIPModel:
+    try:
+        # Check if the weights file exists
+        if not Path(weights_path).exists():
+            logger.info(f"BanglaCLIP weights not found locally, downloading from Hugging Face...")
+            url = "https://huggingface.co/Mansuba/Bangla_text_to_image_app/resolve/main/banglaclip_model_epoch_10.pth"
+            response = requests.get(url, stream=True)
+            response.raise_for_status()
+            
+            # Save the downloaded file
+            with open(weights_path, "wb") as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+            logger.info("BanglaCLIP weights downloaded successfully.")
+
+        # Load the model weights
+        clip_model = CLIPModel.from_pretrained(self.clip_model_name)
+        state_dict = torch.load(weights_path, map_location=self.device)
+
+        cleaned_state_dict = {
+            k.replace('module.', '').replace('clip.', ''): v
+            for k, v in state_dict.items()
+            if k.replace('module.', '').replace('clip.', '').startswith(('text_model.', 'vision_model.'))
+        }
+
+        clip_model.load_state_dict(cleaned_state_dict, strict=False)
+        return clip_model.to(self.device)
+
+    except Exception as e:
+        logger.error(f"Failed to load BanglaCLIP model: {str(e)}")
+        raise
+
+
+
 
 class EnhancedBanglaSDGenerator:
     def __init__(