Update train_model.py

train_model.py

@@ -16,12 +16,11 @@ def train_model(
     hf_token: str
 ):
     try:
-        # Memory and seed setup
         gc.collect()
-        torch.cuda.empty_cache()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
         set_seed(42)
 
-        # Unzip uploaded data
         instance_data_dir = Path("instance_data")
         if instance_data_dir.exists():
             shutil.rmtree(instance_data_dir)
@@ -30,26 +29,24 @@ def train_model(
 
         print(f"✅ Data extracted to: {instance_data_dir}")
 
-        # Load pre-trained model
         model_id = "CompVis/stable-diffusion-v1-4"
         pipe = StableDiffusionPipeline.from_pretrained(
             model_id,
-            torch_dtype=torch.float16,
-            revision="fp16",
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            revision="fp16" if torch.cuda.is_available() else "main",
             use_auth_token=hf_token
         )
-        pipe.to("cuda")
 
-        # Placeholder training simulation
-        print(f"🚧 Starting training for {max_train_steps} steps with LR={learning_rate}...")
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        pipe.to(device)
+
+        print(f"🚧 Simulating training for {max_train_steps} steps at LR={learning_rate}")
         for step in range(int(max_train_steps)):
             if step % 100 == 0 or step == int(max_train_steps) - 1:
-                print(f"Training step {step + 1}/{max_train_steps}")
+                print(f"Step {step + 1}/{max_train_steps}")
 
-        # Save model
         os.makedirs(output_dir, exist_ok=True)
         pipe.save_pretrained(output_dir)
-
         return f"🎉 Training completed. Model saved to: {output_dir}"
 
     except Exception as e: