Update app.py

app.py

@@ -52,7 +52,8 @@ pipe = DiffusionPipeline.from_pretrained(
     kontext_model, 
     quantization_config=quant_config,
     torch_dtype=dtype
-    ).to(device)
+    ).to("cuda")
+torch.cuda.empty_cache()
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
@@ -244,6 +245,7 @@ def infer_with_lora(input_image, prompt, selected_index, custom_lora, seed=42, r
         print(f"Error during inference: {e}")
         return None, seed, gr.update(visible=False)
 
+infer_with_lora.zerogpu = True
 # CSS styling
 css = """
 #main_app {