Update train_model.py

train_model.py

@@ -5,6 +5,16 @@ import gc
 from pathlib import Path
 from diffusers import StableDiffusionPipeline
 from accelerate.utils import set_seed
+from huggingface_hub import HfApi, RepositoryNotFoundError, create_repo
+
+def ensure_repo_exists(repo_id: str, hf_token: str):
+    api = HfApi()
+    try:
+        api.repo_info(repo_id, token=hf_token)
+        print(f"ℹ️ Repo '{repo_id}' already exists.")
+    except RepositoryNotFoundError:
+        create_repo(repo_id=repo_id, token=hf_token, repo_type="dataset", private=False)
+        print(f"✅ Repo '{repo_id}' created.")
 
 def train_model(
     instance_token: str,
@@ -56,6 +66,10 @@ 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("generated-images", hf_token)
+
         return f"🎉 Training completed. Model saved to: {output_dir}"
 
     except Exception as e: