Update train_model.py

train_model.py

@@ -13,13 +13,16 @@ def train_model(
     output_dir: str,
     max_train_steps: int,
     learning_rate: float,
-    hf_token: str
+    hf_token: str,
+    seed: int = 42,
+    precision: str = "fp16"
 ):
     try:
         gc.collect()
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
-        set_seed(42)
+
+        set_seed(seed)
 
         instance_data_dir = Path("instance_data")
         if instance_data_dir.exists():
@@ -30,30 +33,26 @@ def train_model(
         print(f"✅ Data extracted to: {instance_data_dir}")
 
         model_id = "CompVis/stable-diffusion-v1-4"
+        torch_dtype = torch.float16 if precision == "fp16" else torch.float32
+        revision = "fp16" if precision == "fp16" else "main"
+
         pipe = StableDiffusionPipeline.from_pretrained(
             model_id,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            revision="fp16" if torch.cuda.is_available() else "main",
+            torch_dtype=torch_dtype,
+            revision=revision,
             use_auth_token=hf_token
         )
 
         device = "cuda" if torch.cuda.is_available() else "cpu"
         pipe.to(device)
 
-        # This is where real training logic would go in production
         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"Step {step + 1}/{max_train_steps}")
 
-        # Save the model to the desired directory
         os.makedirs(output_dir, exist_ok=True)
         pipe.save_pretrained(output_dir)
-
-        # Create a success flag file
-        with open(os.path.join(output_dir, "training_complete.txt"), "w") as f:
-            f.write("Training complete")
-
         return f"🎉 Training completed. Model saved to: {output_dir}"
 
     except Exception as e: