Update app.py

app.py

@@ -50,8 +50,9 @@ image_examples = [
     
 ]
 
-@spaces.GPU
+@spaces.GPU(enable_queue=True)
 def load_model(base_model_path, lora_path):
+    global pipe 
     transformer = FluxTransformer2DModel.from_pretrained(base_model_path, subfolder='transformer', torch_dtype=torch.bfloat16)
     gr.Info(str(f"Model loading: {int((40 / 100) * 100)}%"))
     # enable image inputs
@@ -77,14 +78,13 @@ def load_model(base_model_path, lora_path):
         base_model_path,
         transformer=transformer,
         torch_dtype=torch.bfloat16
-    )
+    ).to("cuda")
     pipe.transformer.to(torch.bfloat16)
     gr.Info(str(f"Model loading: {int((80 / 100) * 100)}%"))
     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
+@spaces.GPU(enable_queue=True)
 def set_seed(seed):
     torch.manual_seed(seed)
     torch.cuda.manual_seed(seed)
@@ -92,9 +92,8 @@ def set_seed(seed):
     np.random.seed(seed)
     random.seed(seed)
 
-@spaces.GPU  
+@spaces.GPU(enable_queue=True)   
 def predict(
-    pipe,
     input_image,
     prompt,
     ddim_steps,
@@ -148,7 +147,6 @@ def predict(
         gray_image_pil = Image.fromarray(gray_image).convert('L')
     else:
         gray_image_pil = input_image["layers"][0]
-    pipe.to("cuda")
     result = pipe(
         prompt=prompt,
         control_image=input_image["background"].convert("RGB"),
@@ -182,7 +180,6 @@ def predict(
    
 
 def infer(
-    pipe,
     input_image,
     ddim_steps,
     seed,
@@ -192,8 +189,7 @@ def infer(
 ):
     img_path = image_path
     msk_path = mask_path
-    return predict(pipe,
-                   input_image, 
+    return predict(input_image, 
                    removal_prompt, 
                    ddim_steps, 
                    seed,
@@ -279,8 +275,7 @@ with gr.Blocks(
 ) as demo:
     base_model_path = 'black-forest-labs/FLUX.1-dev'
     lora_path = 'theSure/Omnieraser'
-    pipe = None
-    pipe = load_model(base_model_path=base_model_path, lora_path=lora_path)
+    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)
@@ -360,7 +355,6 @@ with gr.Blocks(
     run_button.click(
         fn=infer,
         inputs=[
-            pipe,
             input_image,
             ddim_steps,
             seed,
@@ -370,5 +364,4 @@ with gr.Blocks(
         outputs=[inpaint_result, gallery]
     )
     
-
-demo.launch()
+demo.launch()