Update train_model.py

train_model.py

@@ -11,7 +11,28 @@ try:
 except ImportError as e:
     raise ImportError("huggingface_hub is missing or incompatible. Please ensure it's installed and up to date.") from e
 
-# TODO: Validate version compatibility for diffusers, accelerate, and huggingface_hub
+# Optional: Version compatibility check
+REQUIRED_HF_HUB_VERSION = "0.22.0"
+REQUIRED_DIFFUSERS_VERSION = "0.25.0"
+REQUIRED_ACCELERATE_VERSION = "0.27.2"
+
+def check_versions():
+    import importlib.metadata as metadata
+    try:
+        hf_hub_version = metadata.version("huggingface_hub")
+        diffusers_version = metadata.version("diffusers")
+        accelerate_version = metadata.version("accelerate")
+
+        print(f"🔍 Versions: huggingface_hub={hf_hub_version}, diffusers={diffusers_version}, accelerate={accelerate_version}")
+
+        if hf_hub_version < REQUIRED_HF_HUB_VERSION:
+            raise RuntimeError(f"huggingface_hub must be >= {REQUIRED_HF_HUB_VERSION}")
+        if diffusers_version < REQUIRED_DIFFUSERS_VERSION:
+            raise RuntimeError(f"diffusers must be >= {REQUIRED_DIFFUSERS_VERSION}")
+        if accelerate_version < REQUIRED_ACCELERATE_VERSION:
+            raise RuntimeError(f"accelerate must be >= {REQUIRED_ACCELERATE_VERSION}")
+    except Exception as e:
+        raise RuntimeError(f"❌ Version check failed: {e}")
 
 def ensure_repo_exists(repo_id: str, hf_token: str):
     api = HfApi()
@@ -38,6 +59,8 @@ def train_model(
     dataset_repo_id: str = "generated-images"
 ):
     try:
+        check_versions()
+
         gc.collect()
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
@@ -47,7 +70,6 @@ def train_model(
         set_seed(seed)
         print(f"🔧 Using random seed: {seed}")
 
-        # Validate precision mode
         if precision not in ["fp16", "fp32"]:
             return f"❌ Training failed: Invalid precision mode '{precision}'. Choose 'fp16' or 'fp32'."
 
@@ -81,7 +103,6 @@ def train_model(
         os.makedirs(output_dir, exist_ok=True)
         pipe.save_pretrained(output_dir)
 
-        # Ensure dataset repo exists on Hugging Face
         ensure_repo_exists(dataset_repo_id, hf_token)
 
         return f"🎉 Training completed. Model saved to: {output_dir}"