Update train_model.py

train_model.py

@@ -5,7 +5,7 @@ import gc
 from pathlib import Path
 from diffusers import StableDiffusionPipeline
 from accelerate.utils import set_seed
-from huggingface_hub import HfApi, HfHubHTTPError, create_repo
+from huggingface_hub import HfApi, create_repo
 
 
 def ensure_repo_exists(repo_id: str, hf_token: str):
@@ -13,8 +13,8 @@ def ensure_repo_exists(repo_id: str, hf_token: str):
     try:
         api.repo_info(repo_id, token=hf_token)
         print(f"ℹ️ Repo '{repo_id}' already exists.")
-    except HfHubHTTPError as e:
-        if hasattr(e, 'response') and e.response is not None and e.response.status_code == 404:
+    except Exception as e:
+        if "404" in str(e):
             create_repo(repo_id=repo_id, token=hf_token, repo_type="dataset", private=False)
             print(f"✅ Repo '{repo_id}' created.")
         else: