config updates

Author: Anthony Wu

tools/cloud-replicate-cog/README.md

@@ -3,8 +3,8 @@ This is a quick demo of `mflux` running on Linux CPUs in Replicate's cloud using
 # Development
 
 - `cog build --tag mflux-linux-cpu:test`  # produces image cog-cloud-replicate-cog
-- `cog predict mflux-linux-cpu:test -i prompt=sunset -i steps=1 -i num_outputs=1 -i seed=$RANDOM`
+- `cog predict mflux-linux-cpu:test --setup-timeout 3600 -i prompt=sunset -i steps=1 -i num_outputs=1 -i seed=$RANDOM`
 
 when running the first time, we may need to allow a longer setup timeout for initial model download:
 
-- `cog predict <image> --setup-timeout 3600 ...`
+- `cog predict  mflux-linux-cpu:test --setup-timeout 3600 ...`

tools/cloud-replicate-cog/cog.yaml

@@ -1,3 +1,4 @@
+image: r8.im/anthonywu/mflux-linux-cpu:test
 build:
   cog_runtime: true
   gpu: false

tools/cloud-replicate-cog/generate.py

@@ -8,7 +8,7 @@
 from mflux.config.config import Config
 from mflux.flux.flux import Flux1
 
-MODEL_CACHE = "FLUX.1-schnell"
+MODEL_CACHE = "/data/FLUX.1-schnell"
 MODEL_URL = "https://weights.replicate.delivery/default/black-forest-labs/FLUX.1-schnell/files.tar"
 
 ASPECT_RATIOS = {
@@ -62,9 +62,9 @@ class TextToImage(BasePredictor):
     def setup(self) -> None:
         """Load the model into memory to make running multiple predictions efficient"""
         start = time.time()
-        print("Loading Flux model weights")
-        if not os.path.exists(MODEL_CACHE):
-            download_weights(MODEL_URL, ".")
+        print(f"Loading Flux model weights to {MODEL_CACHE}")
+        if not Path(MODEL_CACHE).exists():
+            download_weights(MODEL_URL, MODEL_CACHE)
         model_path = find_transformers_parent(MODEL_CACHE)
         print(f"model path found at: {model_path}")
         self.flux = Flux1(path=model_path, model="schnell")