laloadrianmorales commited on
Commit
c37c52c
ยท
verified ยท
1 Parent(s): a9779fe

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +36 -104
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"๐Ÿš€ Matrix-Game-2.0 Streamlined")
17
- print(f"๐Ÿ“ฑ Device: {DEVICE}")
18
- print(f"๐Ÿ”ฅ CUDA Available: {torch.cuda.is_available()}")
19
 
20
- # Global variables for model loading
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("๐Ÿ“ฅ Downloading Matrix-Game-2.0 model...")
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"โŒ Setup failed: {e}")
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, "โŒ Please upload an input image first"
72
 
73
- # Setup model if not already done
74
  if not download_and_setup_model():
75
- return None, "โŒ Failed to setup model"
76
 
77
  try:
78
- # Create temporary directories
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
- # Prepare input image
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", quality=95)
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(min(num_frames, 100)),
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
- # Execute with timeout
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', '.gif')):
132
  video_files.append(os.path.join(root, file))
133
 
134
  if video_files:
135
- # Copy to a permanent location
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
- error_log = f"โŒ Generation Failed\n๐Ÿ“ Error output: {process.stderr[:500] if process.stderr else 'No error details'}\n๐Ÿ’ญ Try adjusting parameters or using a different input image"
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"โŒ Error during generation: {str(e)}"
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
- # Simple Gradio Interface
156
- def create_interface():
157
- with gr.Blocks(title="Matrix-Game-2.0") as interface:
158
-
159
- gr.HTML("""
160
- <div style="text-align: center; padding: 30px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color: white; border-radius: 15px; margin-bottom: 20px;">
161
- <h1>๐ŸŽฎ Matrix-Game-2.0</h1>
162
- <p style="font-size: 18px;">Interactive World Model for Real-Time Video Generation</p>
163
- <p style="opacity: 0.8;">Upload an image and generate interactive video content!</p>
164
- </div>
165
- """)
166
-
167
- with gr.Row():
168
- with gr.Column():
169
- gr.Markdown("### ๐Ÿ“ธ Input")
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
- return interface
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)