Convert RGBA inputs to RGB

lucataco · Nov 22, 2023 · d7ffcb2 · d7ffcb2
1 parent 301e77d
commit d7ffcb2
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 1 deletion.
diff --git a/README.md b/README.md
@@ -10,7 +10,7 @@ Follow the [model pushing guide](https://replicate.com/docs/guides/push-a-model)
 
 Download weights first
 
-    cog run python script/download_weights.py
+    cog run script/download-weights
 
 Then for predictions,
 

diff --git a/output.mp4 b/output.mp4
diff --git a/predict.py b/predict.py
@@ -7,6 +7,7 @@
 import yaml
 import torch
 sys.path.extend(['/SEINE'])
+from PIL import Image
 import subprocess
 
 class Predictor(BasePredictor):
@@ -50,6 +51,12 @@ def predict(
         if seed is None:
             seed = int.from_bytes(os.urandom(2), "big")
         print(f"Using seed: {seed}")
+
+        img = Image.open(image)
+        if img.mode == "RGBA":
+            img = img.convert("RGB")
+            image = "/tmp/input.jpg"
+            img.save(image)
 
         # Clean up past runs (just in case)
         os.system("rm -rf /src/results/")