Update app.py

app.py

@@ -84,6 +84,7 @@ def load_model(base_model_path, lora_path):
     gr.Info(str(f"Inject LoRA: {lora_path}"))
     pipe.load_lora_weights(lora_path, weight_name="pytorch_lora_weights.safetensors")
     gr.Info(str(f"Model loading: {int((100 / 100) * 100)}%"))
+    return pipe
 @spaces.GPU
 def set_seed(seed):
     torch.manual_seed(seed)
@@ -147,6 +148,10 @@ def predict(
         gray_image_pil = Image.fromarray(gray_image).convert('L')
     else:
         gray_image_pil = input_image["layers"][0]
+    if pipe is None:
+        base_model_path = 'black-forest-labs/FLUX.1-dev'
+        lora_path = 'theSure/Omnieraser'
+        pipe = load_model(base_model_path, lora_path)
     result = pipe(
         prompt=prompt,
         control_image=input_image["background"].convert("RGB"),
@@ -275,7 +280,7 @@ with gr.Blocks(
 ) as demo:
     base_model_path = 'black-forest-labs/FLUX.1-dev'
     lora_path = 'theSure/Omnieraser'
-    load_model(base_model_path=base_model_path, lora_path=lora_path)
+    a = load_model(base_model_path=base_model_path, lora_path=lora_path)
 
     ddim_steps = gr.Slider(visible=False, value=28)
     scale = gr.Slider(visible=False, value=3.5)