Update app.py

app.py

@@ -4,7 +4,7 @@ import random
 import spaces
 import torch
 import time
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, AutoencoderTiny
 from custom_pipeline import FLUXPipelineWithIntermediateOutputs
 
 # Constants
@@ -19,6 +19,7 @@ dtype = torch.float16
 pipe = FLUXPipelineWithIntermediateOutputs.from_pretrained(
     "black-forest-labs/FLUX.1-schnell", torch_dtype=dtype
 ).to("cuda")
+pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.float16)
 torch.cuda.empty_cache()
 
 # Inference function
@@ -26,7 +27,7 @@ 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(device="cuda").manual_seed(int(float(seed)))
+    generator = torch.Generator().manual_seed(int(float(seed)))
 
     start_time = time.time()