Update app.py

app.py

@@ -16,7 +16,6 @@ with open('loras.json', 'r') as f:
 # Initialize the base model
 base_model = "black-forest-labs/FLUX.1-dev"
 pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
-pipe.to("cuda")
 
 MAX_SEED = 2**32-1
 
@@ -50,6 +49,7 @@ def update_selection(evt: gr.SelectData):
 
 @spaces.GPU(duration=80)
 def generate_image(prompt, trigger_word, steps, seed, cfg_scale, width, height, lora_scale, progress):
+    pipe.to("cuda")
     
     generator = torch.Generator(device="cuda").manual_seed(seed)
     
@@ -87,7 +87,7 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, wid
             seed = random.randint(0, MAX_SEED)
     
     image = generate_image(prompt, trigger_word, steps, seed, cfg_scale, width, height, lora_scale, progress)
-
+    pipe.to("cpu")
     pipe.unload_lora_weights()
     return image, seed