Update app.py

app.py

@@ -26,21 +26,20 @@ torch.cuda.empty_cache()
 def generate_image(prompt, seed=42, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT, randomize_seed=False, num_inference_steps=2, progress=gr.Progress(track_tqdm=True)):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator().manual_seed(int(float(seed)))
+    generator = torch.Generator(device="cuda").manual_seed(int(float(seed)))
 
     start_time = time.time()
 
     # Only generate the last image in the sequence
-    for img in pipe.generate_images(  
+    img = pipe.generate_images( 
             prompt=prompt,
-            guidance_scale=0, # as Flux schnell is guidance free
-            num_inference_steps=num_inference_steps,
             width=width,
             height=height,
+            num_inference_steps=num_inference_steps,
             generator=generator
-        ): 
-        latency = f"Latency: {(time.time()-start_time):.2f} seconds"    
-        yield img, seed, latency
+        )
+    latency = f"Latency: {(time.time()-start_time):.2f} seconds"    
+    return img, seed, latency
 
 # Example prompts
 examples = [