Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
app.py
CHANGED
@@ -1,4 +1,3 @@
|
|
1 |
-
import os
|
2 |
import gradio as gr
|
3 |
import torch
|
4 |
import spaces
|
@@ -6,6 +5,7 @@ from PIL import Image
|
|
6 |
import tempfile
|
7 |
import subprocess
|
8 |
import sys
|
|
|
9 |
from huggingface_hub import snapshot_download
|
10 |
import shutil
|
11 |
|
@@ -13,184 +13,116 @@ import shutil
|
|
13 |
MODEL_REPO = "Skywork/Matrix-Game-2.0"
|
14 |
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
15 |
|
16 |
-
print(f"
|
17 |
-
print(f"
|
18 |
-
print(f"🔥 CUDA Available: {torch.cuda.is_available()}")
|
19 |
|
20 |
-
# Global variables
|
21 |
model_loaded = False
|
22 |
model_path = None
|
23 |
|
24 |
def download_and_setup_model():
|
25 |
-
"""Download model and setup environment - run once"""
|
26 |
global model_loaded, model_path
|
27 |
|
28 |
if model_loaded:
|
29 |
return True
|
30 |
|
31 |
try:
|
32 |
-
print("
|
33 |
-
|
34 |
-
# Download the model to cache
|
35 |
model_path = snapshot_download(
|
36 |
repo_id=MODEL_REPO,
|
37 |
-
cache_dir="./model_cache"
|
38 |
-
allow_patterns=["*.safetensors", "*.bin", "*.json", "*.yaml", "*.yml", "*.py"],
|
39 |
)
|
40 |
|
41 |
-
print(f"✅ Model downloaded to: {model_path}")
|
42 |
-
|
43 |
-
# Clone the inference code repository
|
44 |
if not os.path.exists("Matrix-Game"):
|
45 |
-
print("📥 Cloning Matrix-Game repository...")
|
46 |
result = subprocess.run([
|
47 |
'git', 'clone', 'https://github.com/SkyworkAI/Matrix-Game.git'
|
48 |
], capture_output=True, text=True, timeout=180)
|
49 |
|
50 |
if result.returncode != 0:
|
51 |
-
print(f"❌ Git clone failed: {result.stderr}")
|
52 |
return False
|
53 |
|
54 |
-
# Setup Python path to include Matrix-Game modules
|
55 |
-
matrix_game_path = os.path.join(os.getcwd(), "Matrix-Game", "Matrix-Game-2")
|
56 |
-
if matrix_game_path not in sys.path:
|
57 |
-
sys.path.insert(0, matrix_game_path)
|
58 |
-
|
59 |
model_loaded = True
|
60 |
return True
|
61 |
|
62 |
except Exception as e:
|
63 |
-
print(f"
|
64 |
return False
|
65 |
|
66 |
@spaces.GPU(duration=120)
|
67 |
def generate_video(input_image, num_frames, seed):
|
68 |
-
"""Generate video using Matrix-Game-2.0"""
|
69 |
-
|
70 |
if input_image is None:
|
71 |
-
return None, "
|
72 |
|
73 |
-
# Setup model if not already done
|
74 |
if not download_and_setup_model():
|
75 |
-
return None, "
|
76 |
|
77 |
try:
|
78 |
-
|
79 |
-
temp_dir = tempfile.mkdtemp(prefix="matrix_gen_")
|
80 |
output_dir = os.path.join(temp_dir, "outputs")
|
81 |
os.makedirs(output_dir, exist_ok=True)
|
82 |
|
83 |
-
#
|
84 |
if max(input_image.size) > 512:
|
85 |
ratio = 512 / max(input_image.size)
|
86 |
new_size = (int(input_image.size[0] * ratio), int(input_image.size[1] * ratio))
|
87 |
input_image = input_image.resize(new_size, Image.Resampling.LANCZOS)
|
88 |
|
89 |
input_path = os.path.join(temp_dir, "input.jpg")
|
90 |
-
input_image.save(input_path, "JPEG"
|
91 |
|
92 |
-
# Find the inference script and config
|
93 |
matrix_dir = os.path.join("Matrix-Game", "Matrix-Game-2")
|
94 |
|
95 |
-
# Basic inference command
|
96 |
cmd = [
|
97 |
sys.executable,
|
98 |
os.path.join(matrix_dir, "inference.py"),
|
99 |
"--img_path", input_path,
|
100 |
"--output_folder", output_dir,
|
101 |
-
"--num_output_frames", str(
|
102 |
-
"--seed", str(seed)
|
103 |
]
|
104 |
|
105 |
-
# Add model and config paths if found
|
106 |
-
config_files = []
|
107 |
-
for root, dirs, files in os.walk(matrix_dir):
|
108 |
-
for file in files:
|
109 |
-
if file.endswith(('.yaml', '.yml')) and 'config' in file.lower():
|
110 |
-
config_files.append(os.path.join(root, file))
|
111 |
-
|
112 |
-
if config_files:
|
113 |
-
cmd.extend(["--config_path", config_files[0]])
|
114 |
-
|
115 |
if model_path:
|
116 |
cmd.extend(["--pretrained_model_path", model_path])
|
117 |
|
118 |
-
|
119 |
-
process = subprocess.run(
|
120 |
-
cmd,
|
121 |
-
capture_output=True,
|
122 |
-
text=True,
|
123 |
-
timeout=300,
|
124 |
-
cwd=matrix_dir
|
125 |
-
)
|
126 |
|
127 |
# Find output video
|
128 |
video_files = []
|
129 |
for root, dirs, files in os.walk(output_dir):
|
130 |
for file in files:
|
131 |
-
if file.lower().endswith(('.mp4', '.avi', '.mov'
|
132 |
video_files.append(os.path.join(root, file))
|
133 |
|
134 |
if video_files:
|
135 |
-
|
136 |
-
final_output = f"output_{seed}.mp4"
|
137 |
shutil.copy(video_files[0], final_output)
|
138 |
-
|
139 |
-
log = f"✅ Generation Successful!\n📊 Input: {input_image.size}\n🎬 Frames: {num_frames}\n🎲 Seed: {seed}\n📁 Output: {final_output}"
|
140 |
-
|
141 |
-
return final_output, log
|
142 |
else:
|
143 |
-
|
144 |
-
return None, error_log
|
145 |
|
146 |
-
except subprocess.TimeoutExpired:
|
147 |
-
return None, "❌ Generation timed out (>5 minutes). Try fewer frames."
|
148 |
except Exception as e:
|
149 |
-
return None, f"
|
150 |
finally:
|
151 |
-
# Cleanup
|
152 |
if 'temp_dir' in locals() and os.path.exists(temp_dir):
|
153 |
shutil.rmtree(temp_dir, ignore_errors=True)
|
154 |
|
155 |
-
#
|
156 |
-
|
157 |
-
|
158 |
-
|
159 |
-
|
160 |
-
|
161 |
-
|
162 |
-
|
163 |
-
|
164 |
-
|
165 |
-
|
166 |
-
|
167 |
-
with gr.
|
168 |
-
|
169 |
-
|
170 |
-
input_image = gr.Image(label="Input Image", type="pil")
|
171 |
-
|
172 |
-
gr.Markdown("### ⚙️ Settings")
|
173 |
-
with gr.Row():
|
174 |
-
num_frames = gr.Slider(25, 100, 50, step=25, label="Number of Frames")
|
175 |
-
seed = gr.Number(value=42, label="Seed", precision=0)
|
176 |
-
|
177 |
-
generate_btn = gr.Button("🚀 Generate Video", variant="primary")
|
178 |
-
|
179 |
-
with gr.Column():
|
180 |
-
gr.Markdown("### 🎬 Generated Video")
|
181 |
-
output_video = gr.Video(label="Result")
|
182 |
-
status_log = gr.Textbox(label="Status Log", lines=8)
|
183 |
-
|
184 |
-
# Event handlers
|
185 |
-
generate_btn.click(
|
186 |
-
fn=generate_video,
|
187 |
-
inputs=[input_image, num_frames, seed],
|
188 |
-
outputs=[output_video, status_log]
|
189 |
-
)
|
190 |
|
191 |
-
|
192 |
|
193 |
-
# Launch the app
|
194 |
if __name__ == "__main__":
|
195 |
-
demo = create_interface()
|
196 |
demo.launch(share=True)
|
|
|
|
|
1 |
import gradio as gr
|
2 |
import torch
|
3 |
import spaces
|
|
|
5 |
import tempfile
|
6 |
import subprocess
|
7 |
import sys
|
8 |
+
import os
|
9 |
from huggingface_hub import snapshot_download
|
10 |
import shutil
|
11 |
|
|
|
13 |
MODEL_REPO = "Skywork/Matrix-Game-2.0"
|
14 |
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
15 |
|
16 |
+
print(f"Device: {DEVICE}")
|
17 |
+
print(f"CUDA Available: {torch.cuda.is_available()}")
|
|
|
18 |
|
19 |
+
# Global variables
|
20 |
model_loaded = False
|
21 |
model_path = None
|
22 |
|
23 |
def download_and_setup_model():
|
|
|
24 |
global model_loaded, model_path
|
25 |
|
26 |
if model_loaded:
|
27 |
return True
|
28 |
|
29 |
try:
|
30 |
+
print("Downloading model...")
|
|
|
|
|
31 |
model_path = snapshot_download(
|
32 |
repo_id=MODEL_REPO,
|
33 |
+
cache_dir="./model_cache"
|
|
|
34 |
)
|
35 |
|
|
|
|
|
|
|
36 |
if not os.path.exists("Matrix-Game"):
|
|
|
37 |
result = subprocess.run([
|
38 |
'git', 'clone', 'https://github.com/SkyworkAI/Matrix-Game.git'
|
39 |
], capture_output=True, text=True, timeout=180)
|
40 |
|
41 |
if result.returncode != 0:
|
|
|
42 |
return False
|
43 |
|
|
|
|
|
|
|
|
|
|
|
44 |
model_loaded = True
|
45 |
return True
|
46 |
|
47 |
except Exception as e:
|
48 |
+
print(f"Setup failed: {e}")
|
49 |
return False
|
50 |
|
51 |
@spaces.GPU(duration=120)
|
52 |
def generate_video(input_image, num_frames, seed):
|
|
|
|
|
53 |
if input_image is None:
|
54 |
+
return None, "Please upload an input image first"
|
55 |
|
|
|
56 |
if not download_and_setup_model():
|
57 |
+
return None, "Failed to setup model"
|
58 |
|
59 |
try:
|
60 |
+
temp_dir = tempfile.mkdtemp()
|
|
|
61 |
output_dir = os.path.join(temp_dir, "outputs")
|
62 |
os.makedirs(output_dir, exist_ok=True)
|
63 |
|
64 |
+
# Resize image
|
65 |
if max(input_image.size) > 512:
|
66 |
ratio = 512 / max(input_image.size)
|
67 |
new_size = (int(input_image.size[0] * ratio), int(input_image.size[1] * ratio))
|
68 |
input_image = input_image.resize(new_size, Image.Resampling.LANCZOS)
|
69 |
|
70 |
input_path = os.path.join(temp_dir, "input.jpg")
|
71 |
+
input_image.save(input_path, "JPEG")
|
72 |
|
|
|
73 |
matrix_dir = os.path.join("Matrix-Game", "Matrix-Game-2")
|
74 |
|
|
|
75 |
cmd = [
|
76 |
sys.executable,
|
77 |
os.path.join(matrix_dir, "inference.py"),
|
78 |
"--img_path", input_path,
|
79 |
"--output_folder", output_dir,
|
80 |
+
"--num_output_frames", str(int(num_frames)),
|
81 |
+
"--seed", str(int(seed))
|
82 |
]
|
83 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
84 |
if model_path:
|
85 |
cmd.extend(["--pretrained_model_path", model_path])
|
86 |
|
87 |
+
process = subprocess.run(cmd, capture_output=True, text=True, timeout=300, cwd=matrix_dir)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
88 |
|
89 |
# Find output video
|
90 |
video_files = []
|
91 |
for root, dirs, files in os.walk(output_dir):
|
92 |
for file in files:
|
93 |
+
if file.lower().endswith(('.mp4', '.avi', '.mov')):
|
94 |
video_files.append(os.path.join(root, file))
|
95 |
|
96 |
if video_files:
|
97 |
+
final_output = f"output_{int(seed)}.mp4"
|
|
|
98 |
shutil.copy(video_files[0], final_output)
|
99 |
+
return final_output, f"Success! Generated {int(num_frames)} frames with seed {int(seed)}"
|
|
|
|
|
|
|
100 |
else:
|
101 |
+
return None, f"Generation failed: {process.stderr[:200]}"
|
|
|
102 |
|
|
|
|
|
103 |
except Exception as e:
|
104 |
+
return None, f"Error: {str(e)}"
|
105 |
finally:
|
|
|
106 |
if 'temp_dir' in locals() and os.path.exists(temp_dir):
|
107 |
shutil.rmtree(temp_dir, ignore_errors=True)
|
108 |
|
109 |
+
# Ultra-minimal interface
|
110 |
+
with gr.Blocks() as demo:
|
111 |
+
|
112 |
+
gr.HTML("<h1>Matrix-Game-2.0</h1><p>Interactive World Model</p>")
|
113 |
+
|
114 |
+
with gr.Row():
|
115 |
+
with gr.Column():
|
116 |
+
input_image = gr.Image(type="pil")
|
117 |
+
num_frames = gr.Slider(minimum=25, maximum=100, value=50)
|
118 |
+
seed = gr.Number(value=42)
|
119 |
+
btn = gr.Button("Generate")
|
120 |
+
|
121 |
+
with gr.Column():
|
122 |
+
output_video = gr.Video()
|
123 |
+
status = gr.Textbox()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
124 |
|
125 |
+
btn.click(generate_video, [input_image, num_frames, seed], [output_video, status])
|
126 |
|
|
|
127 |
if __name__ == "__main__":
|
|
|
128 |
demo.launch(share=True)
|