Update app.py

app.py

@@ -4,6 +4,9 @@ import numpy as np
 import random
 import spaces
 import torch
+import tempfile
+import os
+from PIL import Image
 import os
 
 token = os.environ["API_TOKEN"]
@@ -17,9 +20,14 @@ MAX_IMAGE_SIZE = 2048
 def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator().manual_seed(seed)
-    image =  InferenceClient("black-forest-labs/FLUX.1-schnell", token=prompt)
-    return image, seed
+    model =  InferenceClient("black-forest-labs/FLUX.1-schnell", token=prompt)
+    image = model(prompt)
+    temp_dir = tempfile.gettempdir()
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".jpg", dir=temp_dir)
+    image.save(temp_file, format="JPEG")
+    temp_file_path = temp_file.name
+    temp_file.close()
+    return temp_file_path, seed
  
 examples = [
     "a tiny astronaut hatching from an egg on the moon",