Update app.py

app.py

@@ -72,12 +72,24 @@ class ModelManager:
                 )
                 
                 logger.info(f"Loading model {model_name}")
-                model = T5ForConditionalGeneration.from_pretrained(
-                    model_path,
-                    token=HF_TOKEN,
-                    local_files_only=False,
-                    device_map="auto"  # This will handle GPU if available
-                )
+                # Check if accelerate is available
+                try:
+                    import accelerate
+                    logger.info("Accelerate package found, using device_map='auto'")
+                    model = T5ForConditionalGeneration.from_pretrained(
+                        model_path,
+                        token=HF_TOKEN,
+                        local_files_only=False,
+                        device_map="auto"
+                    )
+                except ImportError:
+                    logger.warning("Accelerate package not found, falling back to CPU")
+                    model = T5ForConditionalGeneration.from_pretrained(
+                        model_path,
+                        token=HF_TOKEN,
+                        local_files_only=False
+                    )
+                    model = model.cpu()
                 
                 cls._instances[model_name] = (model, tokenizer)
                 logger.info(f"Successfully loaded {model_name}")