Add download script, fix training tab

.gitignore

@@ -1,7 +1,12 @@
-__pycache
+__pycache__
 .cog
 blip-cache
 blip-proc-cache
 checkpoint
 model-cache
 swin2sr-cache
+temp
+temp_in
+training_out
+output.tar
+trained_model.tar

README.md

@@ -8,7 +8,7 @@ First, download the pre-trained weights:
 
 Then, you can run predictions:
 
-    cog train -i input_images=@zeke.zip -i use_face_detection_instead=True
+    cog predict -i input_images=@zeke.zip -i use_face_detection_instead=True
 
 ## Example:
 

cog.yaml

@@ -21,5 +21,4 @@ build:
     - wget http://thegiflibrary.tumblr.com/post/11565547760 -O face_landmarker_v2_with_blendshapes.task -q https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/1/face_landmarker.task
 
 # predict.py defines how predictions are run on your model
-predict: "predict.py:Predictor"
-train: "train.py:train"
+predict: "predict.py:Predictor"

script/download-weights

@@ -0,0 +1,25 @@
+#!/usr/bin/env python
+
+# Run this before you deploy it on replicate
+import os
+import sys
+import torch
+from diffusers import StableDiffusionXLPipeline
+
+# append project directory to path so predict.py can be imported
+sys.path.append('.')
+from predict import MODEL_NAME, MODEL_CACHE
+
+# Make cache folders
+if not os.path.exists(MODEL_CACHE):
+    os.makedirs(MODEL_CACHE)
+
+# SD-XL-Base-1.0 refiner
+pipe = StableDiffusionXLPipeline.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=torch.float16,
+    use_safetensors=True,
+    variant="fp16"
+)
+
+pipe.save_pretrained(MODEL_CACHE, safe_serialization=True)