Update app.py

app.py

@@ -14,11 +14,11 @@ except ImportError:
     raise ImportError("The 'sentencepiece' library is required but not installed. Please add it to your environment.")
 
 # Set the device and dtype
-dtype = torch.bfloat16
+dtype = torch.float16  # Change to float16 for better compatibility and performance
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Load the diffusion pipeline without requiring an API token
-pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell") 
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
@@ -32,7 +32,7 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_in
     
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator().manual_seed(seed)
+    generator = torch.Generator(device=device).manual_seed(seed)  # Ensure generator is on the correct device
     
     try:
         image = pipe(