Flux bug (#431)

hal9ai · Oct 15, 2024 · d4542af · d4542af
1 parent c8639aa
commit d4542af
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/apps/flux/app.py b/apps/flux/app.py
@@ -1,22 +1,23 @@
 import os
+import base64
+import shutil
 import requests
 import replicate
 
 from PIL import Image
 from io import BytesIO
-import shutil
+
 
 def flux_image(prompt, filename):
   output = replicate.run("black-forest-labs/flux-dev", input={"prompt": prompt})
-  image_url = 'http://' + output[0]
 
-  response = requests.get(image_url)
-  response.raise_for_status()
+  header, encoded = output[0].url.split(',', 1)
 
-  image = Image.open(BytesIO(response.content))
+  data = base64.b64decode(encoded)
+  image = Image.open(BytesIO(data))
   image.save(filename, format="JPEG")
 
   shutil.copy(filename, f"storage/{filename}")
 
 prompt = input()
-flux_image(prompt, "hal9-flux.jpg")
+flux_image(prompt, "hal9-flux.jpg")