fix pipe

app.py

@@ -1,19 +1,20 @@
-import spaces
 import gradio as gr
 import torch
 from PIL import Image
 from diffusers import DiffusionPipeline
-
+import spaces
 
 # Constants
 MAX_SEED = 2**32 - 1
 MAX_IMAGE_SIZE = 2048
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Load FLUX model
-pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.float16)
-pipe.to("cuda")
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
 pipe.enable_model_cpu_offload()
-pipe.enable_vae_slicing()
+pipe.vae.enable_slicing()
+pipe.vae.enable_tiling()
 
 def print_model_shapes(pipe):
     print("Model component shapes:")
@@ -26,7 +27,7 @@ print_model_shapes(pipe)
 
 @spaces.GPU()
 def infer(prompt, init_image=None, seed=None, width=1024, height=1024, num_inference_steps=4, guidance_scale=0.0):
-    generator = torch.Generator(device="cuda").manual_seed(seed) if seed is not None else None
+    generator = torch.Generator(device=device).manual_seed(seed) if seed is not None else None
 
     try:
         if init_image is None: