ghostsInTheMachine
commited on
Commit
•
06f8ae7
1
Parent(s):
12b0993
Update app.py
Browse files
app.py
CHANGED
@@ -10,13 +10,14 @@ from PIL import Image
|
|
10 |
import moviepy.editor as mp
|
11 |
from infer import lotus, load_models
|
12 |
import logging
|
|
|
13 |
|
14 |
# Set up logging
|
15 |
logging.basicConfig(level=logging.INFO)
|
16 |
logger = logging.getLogger(__name__)
|
17 |
|
18 |
-
# Set device to use
|
19 |
-
device = torch.device("cuda
|
20 |
|
21 |
# Load models once
|
22 |
task_name = 'depth'
|
@@ -37,8 +38,6 @@ def preprocess_video(video_path, target_fps=24):
|
|
37 |
def process_frame(frame, seed=0):
|
38 |
"""Process a single frame and return depth map."""
|
39 |
try:
|
40 |
-
torch.cuda.empty_cache() # Clear GPU cache
|
41 |
-
|
42 |
# Convert frame to PIL Image
|
43 |
image = Image.fromarray(frame).convert('RGB')
|
44 |
|
@@ -51,7 +50,7 @@ def process_frame(frame, seed=0):
|
|
51 |
logger.error(f"Error processing frame: {e}")
|
52 |
return None
|
53 |
|
54 |
-
#
|
55 |
def process_video(video_path, fps=0, seed=0):
|
56 |
"""Process video frames individually and generate depth maps."""
|
57 |
# Create a persistent temporary directory
|
|
|
10 |
import moviepy.editor as mp
|
11 |
from infer import lotus, load_models
|
12 |
import logging
|
13 |
+
import spaces # Import spaces for @spaces.GPU decorator
|
14 |
|
15 |
# Set up logging
|
16 |
logging.basicConfig(level=logging.INFO)
|
17 |
logger = logging.getLogger(__name__)
|
18 |
|
19 |
+
# Set device to use GPU if available
|
20 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
21 |
|
22 |
# Load models once
|
23 |
task_name = 'depth'
|
|
|
38 |
def process_frame(frame, seed=0):
|
39 |
"""Process a single frame and return depth map."""
|
40 |
try:
|
|
|
|
|
41 |
# Convert frame to PIL Image
|
42 |
image = Image.fromarray(frame).convert('RGB')
|
43 |
|
|
|
50 |
logger.error(f"Error processing frame: {e}")
|
51 |
return None
|
52 |
|
53 |
+
@spaces.GPU # Decorate the function to use Hugging Face Spaces GPUs
|
54 |
def process_video(video_path, fps=0, seed=0):
|
55 |
"""Process video frames individually and generate depth maps."""
|
56 |
# Create a persistent temporary directory
|