Update app.py

app.py

@@ -50,41 +50,37 @@ image_examples = [
     
 ]
 
-@spaces.GPU
-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
-    with torch.no_grad():
-        initial_input_channels = transformer.config.in_channels
-        new_linear = torch.nn.Linear(
-            transformer.x_embedder.in_features*4,
-            transformer.x_embedder.out_features,
-            bias=transformer.x_embedder.bias is not None,
-            dtype=transformer.dtype,
-            device=transformer.device,
-        )
-        new_linear.weight.zero_()
-        new_linear.weight[:, :initial_input_channels].copy_(transformer.x_embedder.weight)
-
-        if transformer.x_embedder.bias is not None:
-            new_linear.bias.copy_(transformer.x_embedder.bias)
-
-        transformer.x_embedder = new_linear
-        transformer.register_to_config(in_channels=initial_input_channels*4)
-
-    pipe = FluxControlRemovalPipeline.from_pretrained(
-        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
+base_model_path = 'black-forest-labs/FLUX.1-dev'
+lora_path = 'theSure/Omnieraser'  
+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
+with torch.no_grad():
+    initial_input_channels = transformer.config.in_channels
+    new_linear = torch.nn.Linear(
+        transformer.x_embedder.in_features*4,
+        transformer.x_embedder.out_features,
+        bias=transformer.x_embedder.bias is not None,
+        dtype=transformer.dtype,
+        device=transformer.device,
+    )
+    new_linear.weight.zero_()
+    new_linear.weight[:, :initial_input_channels].copy_(transformer.x_embedder.weight)
+    if transformer.x_embedder.bias is not None:
+        new_linear.bias.copy_(transformer.x_embedder.bias)
+    transformer.x_embedder = new_linear
+    transformer.register_to_config(in_channels=initial_input_channels*4)
+pipe = FluxControlRemovalPipeline.from_pretrained(
+    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)}%"))
+
 @spaces.GPU
 def set_seed(seed):
     torch.manual_seed(seed)
@@ -277,9 +273,7 @@ with gr.Blocks(
     ),
     title="Omnieraser"
 ) as demo:
-    base_model_path = 'black-forest-labs/FLUX.1-dev'
-    lora_path = 'theSure/Omnieraser'
-    #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)