Update app.py

app.py

@@ -63,7 +63,7 @@ MODEL_OPTIONS = {
 
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "4096"))
 USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
-ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
+ENABLE_CPU_OFFLOAD = 0
 BATCH_SIZE = int(os.getenv("BATCH_SIZE", "1"))
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -114,11 +114,14 @@ def load_and_prepare_model(model_id):
         model_id,
         torch_dtype=torch.bfloat16,
         add_watermarker=False,
+        use_safetensors=True,
         vae=vae,
     )
     pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
-    #pipe.to(torch.bfloat16)
+    pipe.to(torch.bfloat16)
     pipe.to('cuda')
+    if ENABLE_CPU_OFFLOAD:
+        pipe.enable_model_cpu_offload()
     return pipe
 
 # Preload and compile both models
@@ -182,7 +185,7 @@ def generate(
     global models
     pipe = models[model_choice]
     seed = int(randomize_seed_fn(seed, randomize_seed))
-    generator = torch.Generator(device='cpu').manual_seed(seed)
+    generator = torch.Generator(device='cuda').manual_seed(seed)
 
     prompt, negative_prompt = apply_style(style_selection, prompt, negative_prompt)