Update app.py

app.py

@@ -8,12 +8,13 @@ from PIL import Image
  # os.environ["CUDA_VISIBLE_DEVICES"] = ""
 os.environ["SAFETENSORS_FAST_GPU"] = "1"
 os.putenv("HF_HUB_ENABLE_HF_TRANSFER","1")
-#import torch
-#device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+import torch
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 # Create the gr.State component *outside* the gr.Blocks context
 
-#predictor = gr.State(None)
+predictor = gr.State(None)
+
 def init_predictor(task_type: str):
     from skyreelsinfer import TaskType
     from skyreelsinfer.offload import OffloadConfig
@@ -32,7 +33,7 @@ def init_predictor(task_type: str):
             ),
             use_multiprocessing=False,
         )
-        return "Model loaded successfully!", predictor  # Return predictor
+        return predictor
 
     except (RepositoryNotFoundError, RevisionNotFoundError, EntryNotFoundError) as e:
         return f"Error: Model not found. Details: {e}", None