Update app.py

app.py

@@ -6,7 +6,6 @@ import torch
 import random
 import spaces
 import gradio as gr
-print(gr.__version__)
 import numpy as np
 
 from PIL import Image, ImageCms
@@ -16,8 +15,6 @@ from diffusers.utils import load_image
 from pipeline_flux_control_removal import FluxControlRemovalPipeline
 
 torch.set_grad_enabled(False)
-device = "cuda"
-print(device)
 image_path = mask_path = None
 image_examples = [...] 
 image_path = mask_path =None
@@ -52,6 +49,7 @@ image_examples = [
     ]
     
 ]
+
 @spaces.GPU(duration=120)
 def load_model(base_model_path, lora_path):
     global pipe
@@ -80,13 +78,12 @@ def load_model(base_model_path, lora_path):
         base_model_path,
         transformer=transformer,
         torch_dtype=torch.bfloat16
-    ).to(device)
+    ).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(duration=120)
 def set_seed(seed):
     torch.manual_seed(seed)
@@ -95,7 +92,7 @@ def set_seed(seed):
     np.random.seed(seed)
     random.seed(seed)
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=120)   
 def predict(
     input_image,
     prompt,
@@ -276,7 +273,7 @@ with gr.Blocks(
     ),
     title="Omnieraser"
 ) as demo:
-    base_model_path = "black-forest-labs/FLUX.1-dev"
+    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)
 
@@ -366,6 +363,6 @@ with gr.Blocks(
         ],
         outputs=[inpaint_result, gallery]
     )
-
+    
 
 demo.launch()