fix version model and lowercase model name in train

services/image_generation.py

@@ -19,7 +19,12 @@ def generate_image(model_name, prompt, steps, cfg_scale, width, height, lora_sca
             }
         )
     else:
-        model_name = model_name.lower().replace(' ', '_')
+        # check if the model has a version. the model is something like model='user/model-name:version' but sometimes we just got model='user/model-name' in this case, let get and add the model version
+        if ':' not in model_name:
+            model_version = replicate.models.get(model_name).latest_version.id
+            print(f"Model version: {model_version}")
+            model_name = f"{model_name}:{model_version}"
+            
         img_url = replicate.run(
             model_name,
             input={

services/train_lora.py

@@ -7,6 +7,7 @@ REPLICATE_OWNER = "josebenitezg"
 
 def lora_pipeline(user_id, zip_path, model_name, trigger_word="TOK", steps=1000, lora_rank=16, batch_size=1, autocaption=True, learning_rate=0.0004):
     print(f'Creating dataset for {model_name}')
+    model_name = model_name.lower().replace(' ', '_')
     hf_repo_name = f"joselobenitezg/flux-dev-{model_name}"
     replicate_repo_name = f"josebenitezg/flux-dev-{model_name}"
     create_repo(hf_repo_name, repo_type='model')