Update app.py

app.py

@@ -31,42 +31,13 @@ class ModelCache:
         self.cache_dir = cache_dir
         self.cache_dir.mkdir(parents=True, exist_ok=True)
 
-    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
-
-
-
+    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
 
 class EnhancedBanglaSDGenerator:
     def __init__(
@@ -133,16 +104,11 @@ class EnhancedBanglaSDGenerator:
 
     def _load_banglaclip_model(self, weights_path: str) -> CLIPModel:
         try:
-            model_file_path = Path("model_cache/banglaclip_model_epoch_10.pth")
-            if not model_file_path.exists():
-                logger.info("Downloading BanglaCLIP weights...")
-                torch.hub.download_url_to_file(
-                    'https://huggingface.co/Mansuba/Bangla_text_to_image_app/resolve/main/banglaclip_model_epoch_10.pth',
-                    model_file_path
-                )
+            if not Path(weights_path).exists():
+                raise FileNotFoundError(f"BanglaCLIP weights not found at {weights_path}")
 
             clip_model = CLIPModel.from_pretrained(self.clip_model_name)
-            state_dict = torch.load(model_file_path, map_location=self.device)
+            state_dict = torch.load(weights_path, map_location=self.device)
 
             cleaned_state_dict = {
                 k.replace('module.', '').replace('clip.', ''): v
@@ -273,7 +239,7 @@ def create_gradio_interface():
         nonlocal generator
         if generator is None:
             generator = EnhancedBanglaSDGenerator(
-                banglaclip_weights_path="model_cache/banglaclip_model_epoch_10.pth",
+                banglaclip_weights_path="banglaclip_model_epoch_10_quantized.pth",
                 cache_dir=str(cache_dir)
             )
         return generator
@@ -353,4 +319,5 @@ def create_gradio_interface():
 
 if __name__ == "__main__":
     demo = create_gradio_interface()
+    # Fixed queue configuration for newer Gradio versions
     demo.queue().launch(share=True)