Added: Image Generation & Music Generation -both support for CPU based (debian slim) or gpu (CUDA based) images. -both utilize .sh script that creates and maintains a deticated instruments_venv for dependencies and include heartbeats that prevent terminal passback (api waste). -code_exe tool timeouts have been changes to a rolling window, if idle for 10 seconds it passes back, else there is no max exe time if a process is responsive.

Added: Image Generation & Music Generation -both support for CPU based (debian slim) or gpu (CUDA based) images. -both utilize .sh script that creates and maintains a deticated instruments_venv for dependencies and include heartbeats that prevent terminal passback (api waste). -code_exe tool timeouts have been changes to a rolling window, if idle for 10 seconds it passes back, else there is no max exe time if a process is responsive.

deci committed May 10, 2025 at 16:49 UTC 143135394c27607c1157157efc2889b9ca8e1518
6 files changed +1365
instruments/default/image_generation/image_generation.md new
+20
@@ -0,0 +1,20 @@
1 +# Problem
2 +Generate an image locally using Stable Diffusion
3 +
4 +# Usage (Recommended for All)
5 +Run the wrapper script for maximum compatibility:
6 +```
7 +bash /a0/instruments/default/image_generation/image_generation.sh "<prompt>"
8 +```
9 +- This script will handle all environment setup, venv creation, and dependency installation for both GPU and CPU images.
10 +- You do **not** need to worry about which Python to use or whether the venv exists.
11 +- The generated image will be saved to `/root/generated_images/` with a timestamped filename.
12 +
13 +# Example
14 +```
15 +bash /a0/instruments/default/image_generation/image_generation.sh "a cat under a tree"
16 +```
17 +
18 +# Notes for Automation/Agents
19 +- Always invoke the shell script as shown above.
20 +- Do **not** call the Python script directly; the shell script ensures reliability and compatibility across all environments.
\ No newline at end of file
instruments/default/image_generation/image_generation.py new
+654
@@ -0,0 +1,654 @@
1 +#!/usr/bin/env python3
2 +
3 +import argparse
4 +import os
5 +import sys
6 +import subprocess
7 +import time
8 +from datetime import datetime
9 +import json # Added for parsing pip list output
10 +import threading
11 +
12 +# Define constants
13 +SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
14 +# Old VENV_DIR calculation:
15 +# VENV_DIR = os.path.join(os.path.dirname(SCRIPT_DIR), "..", "instruments_venv")
16 +# New VENV_DIR: Absolute path as defined in Dockerfile.cuda
17 +VENV_DIR = "/opt/instruments_venv"
18 +
19 +DEFAULT_OUTPUT_DIR = "/root/generated_images" # Added back
20 +MODEL_CACHE_DIR = os.path.expanduser("~/.cache/stable-diffusion")
21 +
22 +# Define constants for PyTorch installation
23 +SHOULD_INSTALL_PYTORCH_CUDA_ON_HOST_CUDA = True # True to install CUDA version of PyTorch if host has CUDA - Added back
24 +TARGET_TORCH_PREFIX = "2.6.0" # Major.Minor.Patch, e.g., "2.0.1"
25 +TARGET_TORCH_CUDA_INSTALL_SPEC = "torch==2.6.0+cu124" # Exact spec for CUDA install attempt
26 +
27 +def get_venv_python_executable(venv_dir_path):
28 + """Gets the path to the Python executable in the virtual environment."""
29 + if sys.platform == "win32":
30 + return os.path.join(venv_dir_path, "Scripts", "python.exe")
31 + else:
32 + return os.path.join(venv_dir_path, "bin", "python")
33 +
34 +venv_python_exe = get_venv_python_executable(VENV_DIR) # Initialize globally
35 +
36 +# --- VENV Robustness Debug ---
37 +print("[DEBUG] Current Python:", sys.executable)
38 +print("[DEBUG] Expected venv Python:", venv_python_exe)
39 +if not os.path.exists(venv_python_exe):
40 + print(f"❌ [FATAL] Expected venv Python does not exist: {venv_python_exe}")
41 + sys.exit(1)
42 +if not os.access(venv_python_exe, os.X_OK):
43 + print(f"❌ [FATAL] Expected venv Python is not executable: {venv_python_exe}")
44 + sys.exit(1)
45 +
46 +def check_cuda():
47 + """Check if NVIDIA GPU and CUDA are likely available on the host"""
48 + try:
49 + # Check if nvidia-smi command works
50 + nvidia_smi = subprocess.run(['nvidia-smi'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, timeout=5)
51 + if nvidia_smi.returncode == 0:
52 + print("✅ Host NVIDIA GPU detected via nvidia-smi.")
53 + return True
54 + print("ℹ️ nvidia-smi command failed or returned non-zero. Assuming no NVIDIA GPU for PyTorch CUDA install.")
55 + return False
56 + except FileNotFoundError:
57 + print("ℹ️ nvidia-smi command not found. Assuming no NVIDIA GPU for PyTorch CUDA install.")
58 + return False
59 + except subprocess.TimeoutExpired:
60 + print("⚠️ Timeout running nvidia-smi. Assuming no NVIDIA GPU for PyTorch CUDA install.")
61 + return False
62 + except Exception as e:
63 + print(f"⚠️ Error running nvidia-smi: {e}. Assuming no NVIDIA GPU for PyTorch CUDA install.")
64 + return False
65 +
66 +# Helper function to get installed packages
67 +def get_installed_packages(venv_python_exe):
68 + """Gets a dictionary of installed packages and their versions in the venv."""
69 + cmd = [venv_python_exe, "-m", "pip", "list", "--format=json", "--disable-pip-version-check"]
70 + print(f"🔍 Checking installed packages in instruments venv...") # Updated message
71 + try:
72 + process = subprocess.run(cmd, capture_output=True, text=True, check=True, timeout=60)
73 + installed_list = json.loads(process.stdout)
74 + return {pkg['name'].lower(): pkg['version'] for pkg in installed_list} # Lowercase names
75 + except subprocess.CalledProcessError as e:
76 + print(f"⚠️ Failed to list installed packages. Pip STDERR (truncated): {e.stderr[:500]}")
77 + return {}
78 + except subprocess.TimeoutExpired:
79 + print("⚠️ Timeout while listing installed packages.")
80 + return {}
81 + except json.JSONDecodeError as e:
82 + print(f"⚠️ Failed to parse JSON from pip list: {e}")
83 + return {}
84 + except Exception as e:
85 + print(f"❌ Unexpected error listing packages: {e}")
86 + return {}
87 +
88 +# Helper to verify PyTorch CUDA status in the venv
89 +def verify_venv_pytorch_cuda(venv_python_exe):
90 + """Checks if torch.cuda.is_available() is True in the venv."""
91 + print("🔍 Verifying PyTorch CUDA status in instruments venv (this might take a moment for initial torch import)...") # Updated message
92 + try:
93 + script = "import torch; print(torch.cuda.is_available())"
94 + result = subprocess.run(
95 + [venv_python_exe, "-c", script],
96 + capture_output=True, text=True, check=True, timeout=120 # Increased timeout to 120 seconds
97 + )
98 + available = result.stdout.strip().lower() == "true"
99 + print(f"ℹ️ PyTorch CUDA in instruments venv reports: {'Available' if available else 'Not Available'}") # Updated message
100 + return available
101 + except subprocess.TimeoutExpired:
102 + print(f"⚠️ PyTorch CUDA status check in instruments venv timed out after 120 seconds.") # Updated message
103 + return False
104 + except subprocess.CalledProcessError as e:
105 + print(f"⚠️ Error verifying PyTorch CUDA status in instruments venv: {e.stderr}") # Updated message
106 + return False
107 +
108 +CORE_DEPENDENCIES = { # Ensure this is defined before use in install_requirements
109 + "huggingface_hub": "0.20.3",
110 + "safetensors": "0.4.1",
111 + "accelerate": "0.21.0",
112 + "diffusers": "0.25.0",
113 + "transformers": "4.38.2",
114 + "scipy": "1.15.2" # For diffusers and other potential uses
115 +}
116 +
117 +# Optional, for specific features or performance
118 +XFORMERS_VERSION = "0.0.29.post3"
119 +
120 +def heartbeat_printer(stop_event, message="⏳ Process still running. Monitor terminal for output.", interval=9):
121 + while not stop_event.is_set():
122 + time.sleep(interval)
123 + if not stop_event.is_set():
124 + print(message)
125 +
126 +def install_requirements(venv_python_exe):
127 + """Install required packages into the virtual environment, checking versions first."""
128 + print(f"🔄 Checking/installing dependencies into venv: {VENV_DIR}")
129 +
130 + installed_pkgs = get_installed_packages(venv_python_exe)
131 +
132 + # --- Target Versions Definitions ---
133 + # For PyTorch, the version check is more about the prefix and CUDA capability.
134 + # TARGET_TORCH_VERSION_PREFIX is used to check if a reasonably modern torch is installed.
135 + TARGET_TORCH_PREFIX = "2.6.0" # Major.Minor.Patch, e.g., "2.0.1"
136 + TARGET_TORCH_CUDA_INSTALL_SPEC = "torch==2.6.0+cu124" # Exact spec for CUDA install attempt
137 +
138 + CORE_DEPENDENCIES = {
139 + "huggingface_hub": "0.20.3",
140 + "safetensors": "0.4.1",
141 + "accelerate": "0.21.0"
142 + }
143 + MAIN_DEPENDENCIES = {
144 + "diffusers": "0.25.0",
145 + "transformers": "4.38.2",
146 + "scipy": "1.15.2"
147 + }
148 + XFORMERS_VERSION = "0.0.29.post3"
149 +
150 + def run_pip_command(command_args, action_desc, processing_message_interval=20, overall_timeout=6000):
151 + # Add -v for more verbose pip output and -u for unbuffered Python output for pip itself
152 + cmd = [venv_python_exe, "-u", "-m", "pip", "-v"] + command_args
153 + print(f"🔄 Running: {' '.join(cmd)}")
154 +
155 + process = None # Initialize process variable
156 + try:
157 + # stdout and stderr will go to console by default
158 + process = subprocess.Popen(cmd, text=True, encoding='utf-8', errors='replace')
159 +
160 + start_time = time.time()
161 + last_message_time = start_time
162 +
163 + while True:
164 + current_time = time.time()
165 +
166 + # Check for overall timeout
167 + if current_time - start_time > overall_timeout:
168 + print(f"⚠️ Timeout ({overall_timeout}s) reached for: {action_desc}")
169 + if process:
170 + process.terminate()
171 + try:
172 + process.wait(timeout=5) # Give it a moment to terminate
173 + except subprocess.TimeoutExpired:
174 + print(f"Killing pip process for '{action_desc}' after terminate timeout.")
175 + process.kill()
176 + process.wait() # Wait for kill to complete
177 + print(f"⚠️ Pip process for '{action_desc}' was terminated/killed due to timeout.")
178 + return False
179 +
180 + # Check if process finished
181 + if process:
182 + return_code = process.poll()
183 + if return_code is not None:
184 + if return_code == 0:
185 + print(f"✅ Successfully {action_desc}")
186 + return True
187 + else:
188 + print(f"⚠️ Failed to {action_desc}. Pip process exited with code: {return_code}")
189 + # Pip's own error messages should have already printed to console
190 + return False
191 + else: # Should not happen if Popen succeeds
192 + print(f"❌ Error: Popen process object is None for {action_desc}")
193 + return False
194 + # Print "still processing" message
195 + if current_time - last_message_time > processing_message_interval:
196 + print(f"⏳ Still processing: {action_desc} (running for {int(current_time - start_time)}s)...")
197 + last_message_time = current_time
198 + time.sleep(1) # Poll interval
199 +
200 + except FileNotFoundError:
201 + print(f"❌ Error: The command '{cmd[0]}' was not found. Is Python/pip correctly set up in the venv path?")
202 + return False
203 + except Exception as e:
204 + print(f"❌ Unexpected error during pip process for {action_desc}: {e}")
205 + if process and process.poll() is None: # If process is still running after an unexpected error
206 + print(f"Terminating hanging pip process for '{action_desc}' due to unexpected error.")
207 + process.terminate()
208 + try:
209 + process.wait(timeout=5)
210 + except subprocess.TimeoutExpired:
211 + process.kill()
212 + process.wait()
213 + return False
214 +
215 + # --- PyTorch Bundle Installation Logic ---
216 + host_has_cuda = check_cuda()
217 + pytorch_installed_version = installed_pkgs.get("torch")
218 +
219 + reinstall_pytorch_bundle = False
220 +
221 + if pytorch_installed_version:
222 + print(f"ℹ️ Found existing PyTorch version: {pytorch_installed_version} in instruments venv.") # Updated message
223 + # Check if major.minor matches and if CUDA status is as expected
224 + if not pytorch_installed_version.startswith(TARGET_TORCH_PREFIX.split('.')[0] + '.' + TARGET_TORCH_PREFIX.split('.')[1]):
225 + print(f"⚠️ Existing PyTorch version {pytorch_installed_version} prefix does not match target {TARGET_TORCH_PREFIX}. Will reinstall.")
226 + reinstall_pytorch_bundle = True
227 + else:
228 + # Version prefix matches, now check CUDA status
229 + venv_pytorch_has_cuda = verify_venv_pytorch_cuda(venv_python_exe)
230 + if SHOULD_INSTALL_PYTORCH_CUDA_ON_HOST_CUDA and host_has_cuda and not venv_pytorch_has_cuda:
231 + print(f"⚠️ Host has CUDA, but PyTorch in instruments venv is NOT CUDA-functional. Will reinstall for CUDA.") # Updated message
232 + reinstall_pytorch_bundle = True
233 + elif SHOULD_INSTALL_PYTORCH_CUDA_ON_HOST_CUDA and not host_has_cuda and venv_pytorch_has_cuda:
234 + print(f"⚠️ Host does NOT have CUDA, but PyTorch in instruments venv IS CUDA-functional. Will reinstall for CPU.") # Updated message
235 + reinstall_pytorch_bundle = True
236 + elif not SHOULD_INSTALL_PYTORCH_CUDA_ON_HOST_CUDA and venv_pytorch_has_cuda: # We want CPU, but it has CUDA
237 + print(f"⚠️ PyTorch CUDA installation not desired, but PyTorch in instruments venv IS CUDA-functional. Will reinstall for CPU.") # Updated message
238 + reinstall_pytorch_bundle = True
239 + else:
240 + print(f"✅ Existing PyTorch ({pytorch_installed_version}) in instruments venv meets expectations (CUDA functional: {venv_pytorch_has_cuda}, Host CUDA: {host_has_cuda}).") # Updated message
241 + else:
242 + print(f"ℹ️ PyTorch not found in instruments venv. Will install.") # Updated message
243 + reinstall_pytorch_bundle = True
244 +
245 + if reinstall_pytorch_bundle:
246 + print("🔄 Preparing to install/reinstall PyTorch bundle (torch, torchvision, torchaudio).")
247 +
248 + # Attempt to purge pip cache before critical installations like PyTorch
249 + print("🧹 Attempting to purge pip cache...")
250 + # Use a short timeout for cache purge, it should be quick or fail fast.
251 + run_pip_command(["cache", "purge"], "purged pip cache", overall_timeout=60)
252 +
253 + for pkg_name in ["torch", "torchvision", "torchaudio"]:
254 + if installed_pkgs.get(pkg_name.lower()): # Uninstall if present
255 + run_pip_command(["uninstall", "-y", pkg_name], f"pre-cleaned {pkg_name}")
256 +
257 + if SHOULD_INSTALL_PYTORCH_CUDA_ON_HOST_CUDA and host_has_cuda:
258 + print("✅ Host NVIDIA GPU detected. Attempting to install PyTorch with CUDA support...")
259 + success = run_pip_command(
260 + ["install", TARGET_TORCH_CUDA_INSTALL_SPEC, "torchvision", "torchaudio", "--index-url", "https://download.pytorch.org/whl/cu124"],
261 + f"installed PyTorch with CUDA ({TARGET_TORCH_CUDA_INSTALL_SPEC})"
262 + )
263 + else:
264 + print("ℹ️ Host does not have NVIDIA GPU or CUDA PyTorch install disabled. Installing CPU version of PyTorch...")
265 + success = run_pip_command(
266 + ["install", f"torch=={TARGET_TORCH_PREFIX}", "torchvision", "torchaudio"],
267 + f"installed PyTorch CPU ({TARGET_TORCH_PREFIX})"
268 + )
269 + if success:
270 + # Verify CUDA functionality again after install attempt
271 + venv_pytorch_has_cuda_after_install = verify_venv_pytorch_cuda(venv_python_exe)
272 + if SHOULD_INSTALL_PYTORCH_CUDA_ON_HOST_CUDA and host_has_cuda and not venv_pytorch_has_cuda_after_install:
273 + print(f"⚠️ WARNING: Host has CUDA, but PyTorch in instruments venv is NOT CUDA-functional after installation.") # Updated message
274 + elif SHOULD_INSTALL_PYTORCH_CUDA_ON_HOST_CUDA and host_has_cuda and venv_pytorch_has_cuda_after_install:
275 + print(f"✅ PyTorch in instruments venv is CUDA-functional after installation, as expected.") # Updated message
276 + else:
277 + print(f"❌ Failed to install PyTorch bundle. See pip errors above.")
278 + # Consider if script should exit here or try to continue with other deps
279 +
280 + # --- Install/Verify other core dependencies ---
281 + for dep, version_spec in CORE_DEPENDENCIES.items():
282 + current_version = installed_pkgs.get(dep.lower()) # Ensure consistent key casing
283 +
284 + if current_version == version_spec:
285 + print(f"✅ {dep} ({version_spec}) is already installed and up to date.")
286 + continue
287 +
288 + action = "Installing" if not current_version else f"Updating from {current_version} to"
289 + print(f"🔄 {action} {dep} to {version_spec}.")
290 +
291 + if current_version: # If a version exists but is wrong/different
292 + if not run_pip_command(["uninstall", "-y", dep], f"uninstalling old {dep} ({current_version})"):
293 + print(f"⚠️ Failed to uninstall old {dep}. Attempting to install target version anyway.")
294 +
295 + if not run_pip_command(["install", f"{dep}=={version_spec}"], f"installed {dep}=={version_spec}"):
296 + print(f"❌ Failed to install {dep}=={version_spec}. Aborting dependency installation.")
297 + return False
298 +
299 + # --- Xformers (Conditional) ---
300 + cuda_available_in_venv_pytorch_final = verify_venv_pytorch_cuda(venv_python_exe)
301 + if cuda_available_in_venv_pytorch_final:
302 + current_xformers_version = installed_pkgs.get("xformers")
303 + if current_xformers_version == XFORMERS_VERSION:
304 + print(f"✅ xformers ({XFORMERS_VERSION}) is already installed and up to date.")
305 + else:
306 + action = "Installing" if not current_xformers_version else f"Updating from {current_xformers_version} to"
307 + print(f"🔄 {action} xformers to {XFORMERS_VERSION} for better GPU performance...")
308 + if current_xformers_version:
309 + run_pip_command(["uninstall", "-y", "xformers"], f"uninstalling old xformers ({current_xformers_version})")
310 + if not run_pip_command(["install", f"xformers=={XFORMERS_VERSION}"], f"installed xformers=={XFORMERS_VERSION}"):
311 + print("⚠️ Warning: Failed to install xformers. This is not critical, generation will work without it.")
312 + else:
313 + # If xformers is installed but CUDA is not available, uninstall xformers
314 + if installed_pkgs.get("xformers"):
315 + print("ℹ️ CUDA not available in PyTorch, but xformers is installed. Uninstalling xformers...")
316 + run_pip_command(["uninstall", "-y", "xformers"], "uninstalling xformers (CUDA not available)")
317 + print("ℹ️ Skipping xformers installation as CUDA is not available in the venv's PyTorch.")
318 +
319 + print("✅ Dependency check/installation process complete for venv.")
320 + return True
321 +
322 +def ensure_cpu_dependencies(venv_python_exe):
323 + """Ensure all required CPU dependencies are installed in the venv."""
324 + import subprocess, json
325 + try:
326 + result = subprocess.run([venv_python_exe, "-m", "pip", "list", "--format=json"], capture_output=True, text=True, check=True)
327 + pkgs = {pkg['name'].lower(): pkg['version'] for pkg in json.loads(result.stdout)}
328 + except Exception:
329 + pkgs = {}
330 + if "torch" not in pkgs:
331 + print("🔄 Installing CPU dependencies in venv...")
332 + subprocess.run([venv_python_exe, "-m", "pip", "install", "--upgrade", "pip", "setuptools", "wheel"], check=True)
333 + subprocess.run([
334 + venv_python_exe, "-m", "pip", "install",
335 + "torch==2.6.0", "torchvision", "torchaudio",
336 + "huggingface-hub==0.20.3", "safetensors==0.4.1", "accelerate==0.21.0",
337 + "diffusers==0.25.0", "transformers==4.38.2", "scipy==1.15.2"
338 + ], check=True)
339 + print("✅ CPU dependencies installed.")
340 + else:
341 + print("✅ CPU dependencies already installed in venv.")
342 +
343 +def manage_venv_and_execution():
344 + """Ensures script runs in venv, verifies dependencies for GPU workflow, then re-launches if needed."""
345 + global venv_python_exe # Make sure we update the global if venv is created
346 + venv_python_exe = get_venv_python_executable(VENV_DIR)
347 +
348 + # If already in venv, continue
349 + if sys.executable == venv_python_exe:
350 + print(f"✅ Running in dedicated instruments virtual environment: {VENV_DIR}")
351 + return True
352 +
353 + # Try to detect if CUDA is available (host and torch)
354 + cuda_available = False
355 + try:
356 + import torch
357 + cuda_available = torch.cuda.is_available()
358 + except Exception:
359 + cuda_available = False
360 +
361 + # If CUDA is available, use Dockerfile pre-created venv workflow (existing logic)
362 + if cuda_available:
363 + if not os.path.exists(VENV_DIR):
364 + print(f"❌ [FATAL] Expected venv for GPU workflow does not exist: {VENV_DIR}")
365 + sys.exit(1)
366 + print(f"ℹ️ Instruments virtual environment found at {VENV_DIR}. Verifying dependencies...")
367 + if not install_requirements(venv_python_exe):
368 + print(f"❌ Failed to install/verify requirements in existing instruments venv. Please check errors. Exiting.")
369 + sys.exit(1)
370 + print(f"🔄 Re-launching script with instruments virtual environment Python: {venv_python_exe}")
371 + try:
372 + os.execv(venv_python_exe, [venv_python_exe] + sys.argv)
373 + except Exception as e:
374 + print(f"❌ [FATAL] Failed to re-launch script with venv: {e}")
375 + print(f"👉 Please try activating the venv manually and running the script:")
376 + print(f" {venv_python_exe} {' '.join(sys.argv)}")
377 + sys.exit(1)
378 + print(f"❌ [FATAL] os.execv should not return, but it did. Exiting.")
379 + sys.exit(1)
380 + else:
381 + # For CPU workflow, assume venv and dependencies are already set up by the shell script
382 + print(f"✅ Running in CPU workflow with venv already set up at {VENV_DIR}")
383 + return True
384 +
385 +def print_versions():
386 + """Print installed versions of all relevant packages (expects to run in venv)"""
387 + # This function now assumes it's running inside the venv due to manage_venv_and_execution
388 + print("\n📦 Installed Package Versions (from venv):")
389 + print("-" * 40)
390 +
391 + try:
392 + import torch
393 + print(f"PyTorch: {torch.__version__}")
394 + print(f"CUDA Available (in this PyTorch runtime): {torch.cuda.is_available()}")
395 + if torch.cuda.is_available():
396 + cuda_version = getattr(getattr(torch, 'version', None), 'cuda', None)
397 + print(f"CUDA Version reported by PyTorch: {cuda_version}")
398 + print(f"cuDNN Version: {torch.backends.cudnn.version() if torch.backends.cudnn.is_available() else 'Not available'}")
399 + print(f"GPU: {torch.cuda.get_device_name(0)}")
400 + except ImportError:
401 + print("PyTorch: Not installed or importable in venv")
402 + except Exception as e:
403 + print(f"Error checking PyTorch version: {e}")
404 +
405 + packages = [
406 + "diffusers",
407 + "transformers",
408 + "safetensors",
409 + "accelerate",
410 + "scipy",
411 + "xformers",
412 + "huggingface_hub"
413 + ]
414 +
415 + for package in packages:
416 + try:
417 + module = __import__(package)
418 + version = getattr(module, "__version__", "Unknown version")
419 + print(f"{package}: {version}")
420 + except ImportError:
421 + print(f"{package}: Not installed")
422 +
423 + print("-" * 40)
424 +
425 +def generate_image(prompt, output_dir=DEFAULT_OUTPUT_DIR, seed=None, size=(512, 512)):
426 + """Generate an image using Stable Diffusion
427 +
428 + Args:
429 + prompt: Text prompt describing the image to generate
430 + output_dir: Directory to save the generated image
431 + seed: Random seed for reproducibility (optional)
432 + size: Output image size as (width, height) tuple (default: 512x512)
433 + """
434 + print(f"🖼️ Generating image with prompt: \"{prompt}\"")
435 +
436 + # Dependencies are now handled by manage_venv_and_execution ensuring script runs in venv.
437 +
438 + # Import required libraries (should be from venv)
439 + try:
440 + import torch
441 + from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipeline # Linter might complain, but this is common for diffusers
442 + from PIL import Image
443 + except ImportError as e:
444 + print(f"❌ Critical Error: Failed to import core libraries (torch, diffusers, PIL) from instruments venv: {e}")
445 + print(f"Ensure dependencies were installed correctly in the instruments venv: {VENV_DIR}") # Updated message
446 + sys.exit(1)
447 +
448 + print(f"✅ Using PyTorch {torch.__version__} (from venv)")
449 +
450 + # This check is crucial: it reflects the venv's PyTorch CUDA status
451 + is_cuda_available_runtime = torch.cuda.is_available()
452 + print(f"✅ CUDA available in current PyTorch runtime: {is_cuda_available_runtime}")
453 +
454 + device = "cpu" # Default to CPU
455 + if is_cuda_available_runtime:
456 + device = "cuda"
457 + try:
458 + print(f"✅ Attempting to use CUDA device: {torch.cuda.get_device_name(0)}")
459 + print(f"ℹ️ GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
460 + except Exception as e:
461 + print(f"⚠️ Could not get CUDA device_name or properties, but CUDA is available. Proceeding. Error: {e}")
462 + else:
463 + # If CUDA is not available in PyTorch, double-check with nvidia-smi for user feedback
464 + if check_cuda(): # check_cuda uses nvidia-smi
465 + print("⚠️ PyTorch reports CUDA not available, but nvidia-smi found an NVIDIA GPU.")
466 + print("⚠️ This might indicate a PyTorch installation issue or driver mismatch within the venv.")
467 + print("⚠️ Using CPU (CUDA not available in PyTorch runtime or no NVIDIA GPU detected).")
468 +
469 +
470 + # Set seed if provided
471 + if seed is not None:
472 + torch.manual_seed(seed)
473 + print(f"🎲 Using seed: {seed}")
474 +
475 + # Load the model
476 + print("🔄 Loading Stable Diffusion model...")
477 + start_time = time.time()
478 + # Heartbeat for model loading
479 + model_loading_stop = threading.Event()
480 + model_loading_thread = threading.Thread(target=heartbeat_printer, args=(model_loading_stop,))
481 + model_loading_thread.start()
482 + try:
483 + pipe = StableDiffusionPipeline.from_pretrained(
484 + "stabilityai/stable-diffusion-2-1-base",
485 + torch_dtype=torch.float16 if device == "cuda" else torch.float32, # Use float16 for CUDA
486 + safety_checker=None, # As per original script
487 + cache_dir=MODEL_CACHE_DIR, # Use the global constant
488 + resume_download=True,
489 + use_safetensors=True
490 + )
491 + finally:
492 + model_loading_stop.set()
493 + model_loading_thread.join()
494 +
495 + try:
496 + pipe = pipe.to(device)
497 + print(f"✅ Model successfully moved to {device}.")
498 +
499 + # Enable memory optimizations
500 + pipe.enable_attention_slicing() # Good for both CPU and CUDA
501 +
502 + if device == "cuda":
503 + # Enable xformers if available (it would have been installed if CUDA was primary target)
504 + try:
505 + import xformers # This import is now from the venv
506 + pipe.enable_xformers_memory_efficient_attention()
507 + print("✅ Using xformers for memory efficient attention on CUDA.")
508 + except ImportError:
509 + print("⚠️ xformers not available in venv or import failed, using standard attention on CUDA.")
510 + except Exception as e: # Catch other xformers errors
511 + print(f"⚠️ Error enabling xformers: {e}. Using standard attention.")
512 +
513 + except RuntimeError as e:
514 + if "CUDA" in str(e).upper() and device == "cuda": # Check if error is CUDA related
515 + print(f"⚠️ Error moving model to CUDA: {e}")
516 + print("⚠️ Falling back to CPU for this generation.")
517 + device = "cpu"
518 + pipe = pipe.to(device) # Move to CPU
519 + pipe.enable_attention_slicing() # Ensure attention slicing on CPU too
520 + # If we fell back to CPU, inform user if they have a GPU
521 + if check_cuda():
522 + print("🔄 NOTE: NVIDIA GPU was detected, but an error occurred using CUDA for the model.")
523 + print("🔄 Generation will proceed on CPU. Check PyTorch/CUDA setup in venv if issues persist.")
524 + else:
525 + print(f"❌ Runtime error during model setup or .to(device): {e}")
526 + raise # Re-raise if not a CUDA OOM or similar fallback scenario
527 + except Exception as e: # Catch other .to(device) errors
528 + print(f"❌ Unexpected error during model setup or .to(device): {e}")
529 + raise
530 +
531 +
532 + print(f"✅ Model loaded in {time.time() - start_time:.2f} seconds, configured for {device}")
533 +
534 + # Create output directory if it doesn't exist
535 + os.makedirs(output_dir, exist_ok=True)
536 +
537 + # Generate the image
538 + print(f"🔄 Generating image on {device}...")
539 + gen_start_time = time.time()
540 + # Heartbeat for image generation
541 + gen_stop = threading.Event()
542 + gen_thread = threading.Thread(target=heartbeat_printer, args=(gen_stop,))
543 + gen_thread.start()
544 + try:
545 + with torch.inference_mode():
546 + output = pipe(
547 + prompt=prompt,
548 + num_inference_steps=50 if device == "cuda" else 25, # Adjusted CPU steps
549 + guidance_scale=9,
550 + height=size[1],
551 + width=size[0]
552 + )
553 + finally:
554 + gen_stop.set()
555 + gen_thread.join()
556 +
557 + # Defensive check for output/images
558 + image = None
559 + if isinstance(output, dict) and 'images' in output and isinstance(output['images'], list) and len(output['images']) > 0:
560 + candidate = output['images'][0]
561 + if hasattr(candidate, 'save'):
562 + image = candidate
563 + elif hasattr(output, 'images') and isinstance(output.images, list) and len(output.images) > 0:
564 + candidate = output.images[0]
565 + if hasattr(candidate, 'save'):
566 + image = candidate
567 + elif isinstance(output, (tuple, list)) and len(output) > 0 and hasattr(output[0], 'save'):
568 + image = output[0]
569 + if image is None:
570 + raise RuntimeError("Output from pipeline does not contain an image in the expected format.")
571 +
572 + print(f"✅ Image generated on {device} in {time.time() - gen_start_time:.2f} seconds")
573 +
574 + # Save the image
575 + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
576 + filename = f"image_{timestamp}.png"
577 + filepath = os.path.join(output_dir, filename)
578 +
579 + if hasattr(image, 'save'):
580 + image.save(filepath)
581 + else:
582 + raise RuntimeError("The generated image object does not have a 'save' method. It may not be a PIL.Image.Image.")
583 + print(f"💾 Image saved to: {filepath}")
584 +
585 + return filepath
586 +
587 +def main():
588 + parser = argparse.ArgumentParser(description="Generate images using Stable Diffusion")
589 + parser.add_argument("prompt", nargs="?", type=str, help="The prompt for image generation")
590 + parser.add_argument("--seed", type=int, default=None, help="Random seed for reproducibility")
591 + parser.add_argument("--output-dir", type=str, default=DEFAULT_OUTPUT_DIR, help="Directory to save generated images")
592 + parser.add_argument("--width", type=int, default=512, help="Width of the generated image (default: 512)")
593 + parser.add_argument("--height", type=int, default=512, help="Height of the generated image (default: 512)")
594 + args = parser.parse_args()
595 +
596 + # Check if a prompt was provided
597 + if not args.prompt:
598 + print("❌ No prompt provided. Please specify a prompt.")
599 + print(f"Example: python {os.path.basename(__file__)} 'A majestic dragon'")
600 + return 1
601 +
602 + # Generate the image
603 + filepath = generate_image(
604 + args.prompt,
605 + args.output_dir,
606 + seed=args.seed,
607 + size=(args.width, args.height)
608 + )
609 +
610 + if filepath:
611 + print(f"✨ Image generation completed successfully!")
612 + print_versions() # Print versions after successful generation
613 +
614 + # Add a message if we used CPU but have GPU hardware that PyTorch couldn't use
615 + try:
616 + import torch # Should be venv's torch
617 + if not torch.cuda.is_available() and check_cuda(): # check_cuda for host hardware
618 + print("\n🔄 NOTE: This image was generated on CPU, but an NVIDIA GPU was detected on the host.")
619 + print("🔄 If you intended to use GPU, please check the PyTorch and CUDA driver setup within the virtual environment.")
620 + print(f"🔄 The virtual environment is located at: {VENV_DIR}")
621 + except ImportError:
622 + pass # PyTorch import failed, previous errors would have caught this.
623 + except Exception as e:
624 + print(f"Note: Error during post-generation GPU check: {e}")
625 +
626 + return 0
627 + else:
628 + print("❌ Image generation failed.")
629 + # Check if we installed CUDA support and a GPU is available but CUDA wasn't recognized by PyTorch
630 + try:
631 + import torch
632 + if not torch.cuda.is_available() and check_cuda():
633 + print("\nℹ️ NOTE: An NVIDIA GPU was detected on the host, but PyTorch could not use CUDA.")
634 + print(f"ℹ️ PyTorch (version {torch.__version__}) reported CUDA as unavailable in the current runtime.")
635 + print(f"ℹ️ Dependencies (including PyTorch with CUDA if hardware was detected) were installed into: {VENV_DIR}")
636 + print("ℹ️ Please ensure your NVIDIA drivers are up to date and compatible with the PyTorch CUDA version attempted.")
637 + print("ℹ️ You might need to manually re-trigger dependency installation or debug the venv if issues persist.")
638 + except ImportError:
639 + print("ℹ️ PyTorch is not importable. Dependency installation likely failed.")
640 + except Exception as e:
641 + print(f"Note: Error during failure analysis: {e}")
642 + return 1
643 +
644 +if __name__ == "__main__":
645 + # This block ensures that the script runs inside its dedicated virtual environment.
646 + # If not, it sets up the venv, installs dependencies, and re-launches itself.
647 + if not manage_venv_and_execution():
648 + # This part is reached if execv fails, manage_venv_and_execution will print error and exit.
649 + # However, to be absolutely sure, we can exit here too.
650 + sys.exit(1) # Exit if re-launch failed (though os.execv doesn't return on success)
651 +
652 + # If manage_venv_and_execution() returns True, it means we are already in the venv.
653 + # Or, if it re-launched, the new process starts from here and manage_venv_and_execution() will return True.
654 + sys.exit(main())
\ No newline at end of file
instruments/default/image_generation/image_generation.sh new
+75
@@ -0,0 +1,75 @@
1 +#!/bin/bash
2 +
3 +VENV_DIR="/opt/instruments_venv"
4 +VENV_PY="$VENV_DIR/bin/python"
5 +
6 +echo "==== Starting Image Generation Script ===="
7 +
8 +# Show GPU info if available
9 +if command -v nvidia-smi &> /dev/null; then
10 + echo "✅ NVIDIA GPU detected, displaying information:"
11 + nvidia-smi
12 +
13 + # Get CUDA version
14 + if [ -x "$(command -v nvcc)" ]; then
15 + echo "✅ NVCC (CUDA Compiler) found:"
16 + nvcc --version
17 + else
18 + echo "⚠️ NVCC not found, CUDA development tools may not be installed properly"
19 + fi
20 +
21 + # Check CUDA libraries
22 + echo "Checking CUDA libraries:"
23 + if ldconfig -p | grep -q libcuda.so; then
24 + echo "✅ CUDA libraries found in system path"
25 + ldconfig -p | grep libcuda.so
26 + else
27 + echo "⚠️ CUDA libraries not found in system path"
28 + fi
29 +else
30 + echo "⚠️ No NVIDIA GPU detected (nvidia-smi not found)."
31 +fi
32 +
33 +# Set CUDA env vars if desired
34 +export CUDA_VISIBLE_DEVICES=0
35 +export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128
36 +
37 +# If venv Python does not exist, create venv and install CPU deps
38 +if [ ! -x "$VENV_PY" ]; then
39 + echo "🛠️ venv not found, creating at $VENV_DIR and installing CPU dependencies..."
40 + python3 -m venv "$VENV_DIR"
41 + "$VENV_PY" -m pip install --upgrade pip setuptools wheel
42 + "$VENV_PY" -m pip install \
43 + torch==2.6.0 torchvision torchaudio \
44 + huggingface-hub==0.20.3 safetensors==0.4.1 accelerate==0.21.0 \
45 + diffusers==0.25.0 transformers==4.38.2 scipy==1.15.2
46 + echo "✅ venv created and CPU dependencies installed."
47 +fi
48 +
49 +# Ensure the Python runtime can find CUDA
50 +echo "====== CUDA Environment Variables ======"
51 +echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES"
52 +echo "LD_LIBRARY_PATH: $LD_LIBRARY_PATH"
53 +
54 +echo "====== Running Image Generation ======"
55 +
56 +# Start the Python process in the background
57 +"$VENV_PY" /a0/instruments/default/image_generation/image_generation.py "$@" &
58 +PY_PID=$!
59 +
60 +# Heartbeat loop
61 +while kill -0 $PY_PID 2>/dev/null; do
62 + sleep 9
63 + if kill -0 $PY_PID 2>/dev/null; then
64 + echo "⏳ Process still running. Monitor terminal for output."
65 + fi
66 +done
67 +
68 +wait $PY_PID
69 +status=$?
70 +
71 +if [ $status -eq 0 ]; then
72 + echo "✅ Image generation completed successfully"
73 +else
74 + echo "❌ Image generation failed with error code $status"
75 +fi
\ No newline at end of file
instruments/default/music_generation/music_generation.md new
+39
@@ -0,0 +1,39 @@
1 +# Problem
2 +Generate music locally using MusicGen (transformers) with robust venv, CPU/GPU, and heartbeat support
3 +
4 +# Usage (Recommended for All)
5 +Run the wrapper script for maximum compatibility:
6 +```
7 +bash /a0/instruments/default/music_generation/music_generation.sh "<prompt>"
8 +```
9 +- This script will handle all environment setup, venv creation, and dependency installation for both GPU and CPU workflows.
10 +- You do **not** need to worry about which Python to use or whether the venv exists.
11 +- The generated music will be saved to `/root/generated_music/` with a timestamped filename (WAV format).
12 +
13 +# Example
14 +```
15 +bash /a0/instruments/default/music_generation/music_generation.sh "An upbeat electronic track with a catchy melody"
16 +```
17 +
18 +# Notes for Automation/Agents
19 +- Always invoke the shell script as shown above.
20 +- Do **not** call the Python script directly; the shell script ensures reliability and compatibility across all environments.
21 +- Output files are WAV format by default, saved in `/root/generated_music/`.
22 +
23 +# Options
24 +- `--seed <int>`: Set a random seed for reproducibility
25 +- `--output-dir <path>`: Change the output directory (default: `/root/generated_music`)
26 +- `--duration <seconds>`: Set music duration (if supported by the model)
27 +
28 +# How It Works
29 +1. Checks for NVIDIA GPU and CUDA libraries
30 +2. Sets up a dedicated venv at `/opt/instruments_venv` if needed
31 +3. Installs all required dependencies (PyTorch, transformers, etc.)
32 +4. Runs the music generation Python script with heartbeat monitoring
33 +5. Outputs a WAV file in `/root/generated_music/` with a timestamped filename
34 +
35 +# Troubleshooting
36 +- If you have a GPU but music is generated on CPU, check the terminal output for CUDA/PyTorch warnings.
37 +- If you see dependency errors, try deleting `/opt/instruments_venv` and rerunning the script.
38 +- The first run may take several minutes to download models and set up the environment.
39 +- For best results, use clear, descriptive prompts (e.g., "A relaxing piano melody with gentle strings").
\ No newline at end of file
instruments/default/music_generation/music_generation.py new
+494
@@ -0,0 +1,494 @@
1 +#!/usr/bin/env python3
2 +
3 +import argparse
4 +import os
5 +import sys
6 +import subprocess
7 +import time
8 +from datetime import datetime
9 +import json
10 +import threading
11 +import torch
12 +
13 +# Define constants
14 +SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
15 +VENV_DIR = "/opt/instruments_venv"
16 +DEFAULT_OUTPUT_DIR = "/root/generated_music"
17 +MODEL_CACHE_DIR = os.path.expanduser("~/.cache/audiocraft-models")
18 +
19 +# PyTorch/CUDA logic
20 +SHOULD_INSTALL_PYTORCH_CUDA_ON_HOST_CUDA = True
21 +TARGET_TORCH_PREFIX = "2.6.0"
22 +TARGET_TORCH_CUDA_INSTALL_SPEC = "torch==2.6.0+cu124"
23 +
24 +# Helper: get venv python
25 +
26 +def get_venv_python_executable(venv_dir_path):
27 + if sys.platform == "win32":
28 + return os.path.join(venv_dir_path, "Scripts", "python.exe")
29 + else:
30 + return os.path.join(venv_dir_path, "bin", "python")
31 +
32 +venv_python_exe = get_venv_python_executable(VENV_DIR)
33 +
34 +# --- VENV Robustness Debug ---
35 +print("[DEBUG] Current Python:", sys.executable)
36 +print("[DEBUG] Expected venv Python:", venv_python_exe)
37 +if not os.path.exists(venv_python_exe):
38 + print(f"❌ [FATAL] Expected venv Python does not exist: {venv_python_exe}")
39 + sys.exit(1)
40 +if not os.access(venv_python_exe, os.X_OK):
41 + print(f"❌ [FATAL] Expected venv Python is not executable: {venv_python_exe}")
42 + sys.exit(1)
43 +
44 +def check_cuda():
45 + try:
46 + nvidia_smi = subprocess.run(['nvidia-smi'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, timeout=5)
47 + if nvidia_smi.returncode == 0:
48 + print("✅ Host NVIDIA GPU detected via nvidia-smi.")
49 + return True
50 + print("ℹ️ nvidia-smi command failed or returned non-zero. Assuming no NVIDIA GPU for PyTorch CUDA install.")
51 + return False
52 + except FileNotFoundError:
53 + print("ℹ️ nvidia-smi command not found. Assuming no NVIDIA GPU for PyTorch CUDA install.")
54 + return False
55 + except subprocess.TimeoutExpired:
56 + print("⚠️ Timeout running nvidia-smi. Assuming no NVIDIA GPU for PyTorch CUDA install.")
57 + return False
58 + except Exception as e:
59 + print(f"⚠️ Error running nvidia-smi: {e}. Assuming no NVIDIA GPU for PyTorch CUDA install.")
60 + return False
61 +
62 +def get_installed_packages(venv_python_exe):
63 + cmd = [venv_python_exe, "-m", "pip", "list", "--format=json", "--disable-pip-version-check"]
64 + print(f"🔍 Checking installed packages in instruments venv...")
65 + try:
66 + process = subprocess.run(cmd, capture_output=True, text=True, check=True, timeout=60)
67 + installed_list = json.loads(process.stdout)
68 + return {pkg['name'].lower(): pkg['version'] for pkg in installed_list}
69 + except Exception as e:
70 + print(f"❌ Unexpected error listing packages: {e}")
71 + return {}
72 +
73 +def verify_venv_pytorch_cuda(venv_python_exe):
74 + print("🔍 Verifying PyTorch CUDA status in instruments venv (this might take a moment for initial torch import)...")
75 + try:
76 + script = "import torch; print(torch.cuda.is_available())"
77 + result = subprocess.run(
78 + [venv_python_exe, "-c", script],
79 + capture_output=True, text=True, check=True, timeout=120
80 + )
81 + available = result.stdout.strip().lower() == "true"
82 + print(f"ℹ️ PyTorch CUDA in instruments venv reports: {'Available' if available else 'Not Available'}")
83 + return available
84 + except Exception as e:
85 + print(f"⚠️ Error verifying PyTorch CUDA status in instruments venv: {e}")
86 + return False
87 +
88 +CORE_DEPENDENCIES = {
89 + "huggingface_hub": "0.20.3",
90 + "safetensors": "0.4.1",
91 + "accelerate": "0.21.0",
92 + "transformers": "4.38.2",
93 + "einops": "0.6.1",
94 + "tqdm": "4.65.0",
95 + "librosa": "0.10.0.post2",
96 + "scipy": "1.12.0",
97 + "numpy": "1.24.3"
98 +}
99 +XFORMERS_VERSION = "0.0.29.post3"
100 +
101 +# Heartbeat printer
102 +def heartbeat_printer(stop_event, message="⏳ Process still running. Monitor terminal for output.", interval=9):
103 + while not stop_event.is_set():
104 + time.sleep(interval)
105 + if not stop_event.is_set():
106 + print(message)
107 +
108 +def install_requirements(venv_python_exe):
109 + print(f"🔄 Checking/installing dependencies into venv: {VENV_DIR}")
110 + installed_pkgs = get_installed_packages(venv_python_exe)
111 + TARGET_TORCH_PREFIX = "2.6.0"
112 + TARGET_TORCH_CUDA_INSTALL_SPEC = "torch==2.6.0+cu124"
113 + def run_pip_command(command_args, action_desc, processing_message_interval=20, overall_timeout=6000):
114 + cmd = [venv_python_exe, "-u", "-m", "pip", "-v"] + command_args
115 + print(f"🔄 Running: {' '.join(cmd)}")
116 + process = None
117 + try:
118 + process = subprocess.Popen(cmd, text=True, encoding='utf-8', errors='replace')
119 + start_time = time.time()
120 + last_message_time = start_time
121 + while True:
122 + current_time = time.time()
123 + if current_time - start_time > overall_timeout:
124 + print(f"⚠️ Timeout ({overall_timeout}s) reached for: {action_desc}")
125 + if process:
126 + process.terminate()
127 + try:
128 + process.wait(timeout=5)
129 + except subprocess.TimeoutExpired:
130 + print(f"Killing pip process for '{action_desc}' after terminate timeout.")
131 + process.kill()
132 + process.wait()
133 + print(f"⚠️ Pip process for '{action_desc}' was terminated/killed due to timeout.")
134 + return False
135 + if process:
136 + return_code = process.poll()
137 + if return_code is not None:
138 + if return_code == 0:
139 + print(f"✅ Successfully {action_desc}")
140 + return True
141 + else:
142 + print(f"⚠️ Failed to {action_desc}. Pip process exited with code: {return_code}")
143 + return False
144 + else:
145 + print(f"❌ Error: Popen process object is None for {action_desc}")
146 + return False
147 + if current_time - last_message_time > processing_message_interval:
148 + print(f"⏳ Still processing: {action_desc} (running for {int(current_time - start_time)}s)...")
149 + last_message_time = current_time
150 + time.sleep(1)
151 + except FileNotFoundError:
152 + print(f"❌ Error: The command '{cmd[0]}' was not found. Is Python/pip correctly set up in the venv path?")
153 + return False
154 + except Exception as e:
155 + print(f"❌ Unexpected error during pip process for {action_desc}: {e}")
156 + if process and process.poll() is None:
157 + print(f"Terminating hanging pip process for '{action_desc}' due to unexpected error.")
158 + process.terminate()
159 + try:
160 + process.wait(timeout=5)
161 + except subprocess.TimeoutExpired:
162 + process.kill()
163 + process.wait()
164 + return False
165 + # --- PyTorch Bundle Installation Logic ---
166 + host_has_cuda = check_cuda()
167 + pytorch_installed_version = installed_pkgs.get("torch")
168 + reinstall_pytorch_bundle = False
169 + if pytorch_installed_version:
170 + print(f"ℹ️ Found existing PyTorch version: {pytorch_installed_version} in instruments venv.")
171 + if not pytorch_installed_version.startswith(TARGET_TORCH_PREFIX.split('.')[0] + '.' + TARGET_TORCH_PREFIX.split('.')[1]):
172 + print(f"⚠️ Existing PyTorch version {pytorch_installed_version} prefix does not match target {TARGET_TORCH_PREFIX}. Will reinstall.")
173 + reinstall_pytorch_bundle = True
174 + else:
175 + venv_pytorch_has_cuda = verify_venv_pytorch_cuda(venv_python_exe)
176 + if SHOULD_INSTALL_PYTORCH_CUDA_ON_HOST_CUDA and host_has_cuda and not venv_pytorch_has_cuda:
177 + print(f"⚠️ Host has CUDA, but PyTorch in instruments venv is NOT CUDA-functional. Will reinstall for CUDA.")
178 + reinstall_pytorch_bundle = True
179 + elif SHOULD_INSTALL_PYTORCH_CUDA_ON_HOST_CUDA and not host_has_cuda and venv_pytorch_has_cuda:
180 + print(f"⚠️ Host does NOT have CUDA, but PyTorch in instruments venv IS CUDA-functional. Will reinstall for CPU.")
181 + reinstall_pytorch_bundle = True
182 + elif not SHOULD_INSTALL_PYTORCH_CUDA_ON_HOST_CUDA and venv_pytorch_has_cuda:
183 + print(f"⚠️ PyTorch CUDA installation not desired, but PyTorch in instruments venv IS CUDA-functional. Will reinstall for CPU.")
184 + reinstall_pytorch_bundle = True
185 + else:
186 + print(f"✅ Existing PyTorch ({pytorch_installed_version}) in instruments venv meets expectations (CUDA functional: {venv_pytorch_has_cuda}, Host CUDA: {host_has_cuda}).")
187 + else:
188 + print(f"ℹ️ PyTorch not found in instruments venv. Will install.")
189 + reinstall_pytorch_bundle = True
190 + if reinstall_pytorch_bundle:
191 + print("🔄 Preparing to install/reinstall PyTorch bundle (torch, torchvision, torchaudio).")
192 + print("🧹 Attempting to purge pip cache...")
193 + run_pip_command(["cache", "purge"], "purged pip cache", overall_timeout=60)
194 + for pkg_name in ["torch", "torchvision", "torchaudio"]:
195 + if installed_pkgs.get(pkg_name.lower()):
196 + run_pip_command(["uninstall", "-y", pkg_name], f"pre-cleaned {pkg_name}")
197 + if SHOULD_INSTALL_PYTORCH_CUDA_ON_HOST_CUDA and host_has_cuda:
198 + print("✅ Host NVIDIA GPU detected. Attempting to install PyTorch with CUDA support...")
199 + success = run_pip_command(
200 + ["install", TARGET_TORCH_CUDA_INSTALL_SPEC, "torchvision", "torchaudio", "--index-url", "https://download.pytorch.org/whl/cu124"],
201 + f"installed PyTorch with CUDA ({TARGET_TORCH_CUDA_INSTALL_SPEC})"
202 + )
203 + else:
204 + print("ℹ️ Host does not have NVIDIA GPU or CUDA PyTorch install disabled. Installing CPU version of PyTorch...")
205 + success = run_pip_command(
206 + ["install", f"torch=={TARGET_TORCH_PREFIX}", "torchvision", "torchaudio"],
207 + f"installed PyTorch CPU ({TARGET_TORCH_PREFIX})"
208 + )
209 + if success:
210 + venv_pytorch_has_cuda_after_install = verify_venv_pytorch_cuda(venv_python_exe)
211 + if SHOULD_INSTALL_PYTORCH_CUDA_ON_HOST_CUDA and host_has_cuda and not venv_pytorch_has_cuda_after_install:
212 + print(f"⚠️ WARNING: Host has CUDA, but PyTorch in instruments venv is NOT CUDA-functional after installation.")
213 + elif SHOULD_INSTALL_PYTORCH_CUDA_ON_HOST_CUDA and host_has_cuda and venv_pytorch_has_cuda_after_install:
214 + print(f"✅ PyTorch in instruments venv is CUDA-functional after installation, as expected.")
215 + else:
216 + print(f"❌ Failed to install PyTorch bundle. See pip errors above.")
217 + # --- Install/Verify other core dependencies ---
218 + for dep, version_spec in CORE_DEPENDENCIES.items():
219 + current_version = installed_pkgs.get(dep.lower())
220 + if current_version == version_spec:
221 + print(f"✅ {dep} ({version_spec}) is already installed and up to date.")
222 + continue
223 + action = "Installing" if not current_version else f"Updating from {current_version} to"
224 + print(f"🔄 {action} {dep} to {version_spec}.")
225 + if current_version:
226 + if not run_pip_command(["uninstall", "-y", dep], f"uninstalling old {dep} ({current_version})"):
227 + print(f"⚠️ Failed to uninstall old {dep}. Attempting to install target version anyway.")
228 + if not run_pip_command(["install", f"{dep}=={version_spec}"], f"installed {dep}=={version_spec}"):
229 + print(f"❌ Failed to install {dep}=={version_spec}. Aborting dependency installation.")
230 + return False
231 + # --- Xformers (Conditional) ---
232 + cuda_available_in_venv_pytorch_final = verify_venv_pytorch_cuda(venv_python_exe)
233 + if cuda_available_in_venv_pytorch_final:
234 + current_xformers_version = installed_pkgs.get("xformers")
235 + if current_xformers_version == XFORMERS_VERSION:
236 + print(f"✅ xformers ({XFORMERS_VERSION}) is already installed and up to date.")
237 + else:
238 + action = "Installing" if not current_xformers_version else f"Updating from {current_xformers_version} to"
239 + print(f"🔄 {action} xformers to {XFORMERS_VERSION} for better GPU performance...")
240 + if current_xformers_version:
241 + run_pip_command(["uninstall", "-y", "xformers"], f"uninstalling old xformers ({current_xformers_version})")
242 + if not run_pip_command(["install", f"xformers=={XFORMERS_VERSION}"], f"installed xformers=={XFORMERS_VERSION}"):
243 + print("⚠️ Warning: Failed to install xformers. This is not critical, generation will work without it.")
244 + else:
245 + if installed_pkgs.get("xformers"):
246 + print("ℹ️ CUDA not available in PyTorch, but xformers is installed. Uninstalling xformers...")
247 + run_pip_command(["uninstall", "-y", "xformers"], "uninstalling xformers (CUDA not available)")
248 + print("ℹ️ Skipping xformers installation as CUDA is not available in the venv's PyTorch.")
249 + print("✅ Dependency check/installation process complete for venv.")
250 + return True
251 +
252 +def manage_venv_and_execution():
253 + global venv_python_exe
254 + venv_python_exe = get_venv_python_executable(VENV_DIR)
255 + if sys.executable == venv_python_exe:
256 + print(f"✅ Running in dedicated instruments virtual environment: {VENV_DIR}")
257 + return True
258 + cuda_available = False
259 + try:
260 + import torch
261 + cuda_available = torch.cuda.is_available()
262 + except Exception:
263 + cuda_available = False
264 + if cuda_available:
265 + if not os.path.exists(VENV_DIR):
266 + print(f"❌ [FATAL] Expected venv for GPU workflow does not exist: {VENV_DIR}")
267 + sys.exit(1)
268 + print(f"ℹ️ Instruments virtual environment found at {VENV_DIR}. Verifying dependencies...")
269 + if not install_requirements(venv_python_exe):
270 + print(f"❌ Failed to install/verify requirements in existing instruments venv. Please check errors. Exiting.")
271 + sys.exit(1)
272 + print(f"🔄 Re-launching script with instruments virtual environment Python: {venv_python_exe}")
273 + try:
274 + os.execv(venv_python_exe, [venv_python_exe] + sys.argv)
275 + except Exception as e:
276 + print(f"❌ [FATAL] Failed to re-launch script with venv: {e}")
277 + print(f"👉 Please try activating the venv manually and running the script:")
278 + print(f" {venv_python_exe} {' '.join(sys.argv)}")
279 + sys.exit(1)
280 + print(f"❌ [FATAL] os.execv should not return, but it did. Exiting.")
281 + sys.exit(1)
282 + else:
283 + print(f"✅ Running in CPU workflow with venv already set up at {VENV_DIR}")
284 + return True
285 +
286 +def print_versions():
287 + print("\n📦 Installed Package Versions (from venv):")
288 + print("-" * 40)
289 + try:
290 + import torch
291 + print(f"PyTorch: {torch.__version__}")
292 + print(f"CUDA Available (in this PyTorch runtime): {torch.cuda.is_available()}")
293 + if torch.cuda.is_available():
294 + cuda_version = getattr(getattr(torch, 'version', None), 'cuda', None)
295 + print(f"CUDA Version reported by PyTorch: {cuda_version}")
296 + print(f"cuDNN Version: {torch.backends.cudnn.version() if torch.backends.cudnn.is_available() else 'Not available'}")
297 + print(f"GPU: {torch.cuda.get_device_name(0)}")
298 + except ImportError:
299 + print("PyTorch: Not installed or importable in venv")
300 + except Exception as e:
301 + print(f"Error checking PyTorch version: {e}")
302 + packages = [
303 + "transformers",
304 + "safetensors",
305 + "accelerate",
306 + "scipy",
307 + "xformers",
308 + "huggingface_hub",
309 + "librosa",
310 + "einops",
311 + "numpy"
312 + ]
313 + for package in packages:
314 + try:
315 + module = __import__(package)
316 + version = getattr(module, "__version__", "Unknown version")
317 + print(f"{package}: {version}")
318 + except ImportError:
319 + print(f"{package}: Not installed")
320 + print("-" * 40)
321 +
322 +def generate_music(prompt, output_dir=DEFAULT_OUTPUT_DIR, duration=None, seed=None):
323 + print(f"🎵 Generating music with prompt: \"{prompt}\"")
324 + try:
325 + import torch
326 + from transformers import AutoProcessor, MusicgenForConditionalGeneration
327 + import scipy.io.wavfile as wavfile
328 + import numpy as np
329 + import subprocess
330 + except ImportError as e:
331 + print(f"❌ Critical Error: Failed to import core libraries (torch, transformers, scipy, numpy) from instruments venv: {e}")
332 + print(f"Ensure dependencies were installed correctly in the instruments venv: {VENV_DIR}")
333 + sys.exit(1)
334 + print(f"✅ Using PyTorch {torch.__version__} (from venv)")
335 + is_cuda_available_runtime = torch.cuda.is_available()
336 + print(f"✅ CUDA available in current PyTorch runtime: {is_cuda_available_runtime}")
337 + device = "cpu"
338 + if is_cuda_available_runtime:
339 + device = "cuda"
340 + try:
341 + print(f"✅ Attempting to use CUDA device: {torch.cuda.get_device_name(0)}")
342 + print(f"ℹ️ GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
343 + except Exception as e:
344 + print(f"⚠️ Could not get CUDA device_name or properties, but CUDA is available. Proceeding. Error: {e}")
345 + else:
346 + if check_cuda():
347 + print("⚠️ PyTorch reports CUDA not available, but nvidia-smi found an NVIDIA GPU.")
348 + print("⚠️ This might indicate a PyTorch installation issue or driver mismatch within the venv.")
349 + print("⚠️ Using CPU (CUDA not available in PyTorch runtime or no NVIDIA GPU detected).")
350 + if seed is not None:
351 + torch.manual_seed(seed)
352 + print(f"🎲 Using seed: {seed}")
353 + print("🔄 Loading MusicGen model...")
354 + start_time = time.time()
355 + model_loading_stop = threading.Event()
356 + model_loading_thread = threading.Thread(target=heartbeat_printer, args=(model_loading_stop,))
357 + model_loading_thread.start()
358 + try:
359 + processor = AutoProcessor.from_pretrained("facebook/musicgen-small", cache_dir=MODEL_CACHE_DIR)
360 + model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small", cache_dir=MODEL_CACHE_DIR)
361 + finally:
362 + model_loading_stop.set()
363 + model_loading_thread.join()
364 + try:
365 + model.to(torch.device(device))
366 + print(f"✅ Model successfully moved to {device}.")
367 + except RuntimeError as e:
368 + if "CUDA" in str(e).upper() and device == "cuda":
369 + print(f"⚠️ Error moving model to CUDA: {e}")
370 + print("⚠️ Falling back to CPU for this generation.")
371 + device = "cpu"
372 + model.to(torch.device(device))
373 + else:
374 + print(f"❌ Runtime error during model setup or .to(device): {e}")
375 + raise
376 + except Exception as e:
377 + print(f"❌ Unexpected error during model setup or .to(device): {e}")
378 + raise
379 + print(f"✅ Model loaded in {time.time() - start_time:.2f} seconds, configured for {device}")
380 + os.makedirs(output_dir, exist_ok=True)
381 + print(f"🔄 Generating music on {device}...")
382 + gen_start_time = time.time()
383 + gen_stop = threading.Event()
384 + gen_thread = threading.Thread(target=heartbeat_printer, args=(gen_stop,))
385 + gen_thread.start()
386 + try:
387 + inputs = processor(
388 + text=[prompt],
389 + padding=True,
390 + return_tensors="pt",
391 + )
392 + # Move each tensor in the BatchEncoding to the correct device
393 + inputs = {k: v.to(torch.device(device)) if hasattr(v, 'to') else v for k, v in inputs.items()}
394 +
395 + # Determine max_new_tokens allowed by the model
396 + max_model_tokens = getattr(model.config, 'max_position_embeddings', None)
397 + if max_model_tokens is None:
398 + # Try to get from model.config.audio_encoder if available
399 + audio_encoder = getattr(model.config, 'audio_encoder', None)
400 + max_model_tokens = getattr(audio_encoder, 'max_position_embeddings', None)
401 + if max_model_tokens is None:
402 + max_model_tokens = 1024 # Safe fallback
403 +
404 + # Calculate max_new_tokens from duration if provided
405 + max_new_tokens = max_model_tokens
406 + if duration is not None:
407 + frame_rate = None
408 + audio_encoder = getattr(model.config, 'audio_encoder', None)
409 + if audio_encoder is not None and hasattr(audio_encoder, 'frame_rate'):
410 + frame_rate = audio_encoder.frame_rate
411 + elif hasattr(model.config, 'frame_rate'):
412 + frame_rate = model.config.frame_rate
413 + if frame_rate is not None:
414 + requested_tokens = int(duration * frame_rate)
415 + if requested_tokens > max_model_tokens:
416 + print(f"⚠️ Requested duration ({duration}s) exceeds model's max token capacity. Clamping to {max_model_tokens / frame_rate:.2f} seconds.")
417 + max_new_tokens = min(requested_tokens, max_model_tokens)
418 + else:
419 + print("⚠️ Could not determine model frame rate. Using model's max token capacity.")
420 + else:
421 + print(f"ℹ️ No duration specified. Using model's max token capacity: {max_model_tokens} tokens.")
422 +
423 + with torch.inference_mode():
424 + audio_values = model.generate(**inputs, max_new_tokens=max_new_tokens)
425 + finally:
426 + gen_stop.set()
427 + gen_thread.join()
428 + # Robustly extract sampling_rate
429 + sampling_rate = 32000 # Default fallback
430 + audio_encoder = getattr(model.config, 'audio_encoder', None)
431 + if audio_encoder is not None and hasattr(audio_encoder, 'sampling_rate'):
432 + sampling_rate = audio_encoder.sampling_rate
433 + elif hasattr(model.config, 'sampling_rate'):
434 + sampling_rate = model.config.sampling_rate
435 + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
436 + filename = f"music_{timestamp}.wav"
437 + filepath = os.path.join(output_dir, filename)
438 + audio_data = audio_values[0, 0].cpu().numpy()
439 + wavfile.write(filepath, sampling_rate, audio_data)
440 + print(f"💾 Music saved to: {filepath}")
441 + print(f"✅ Music generated on {device} in {time.time() - gen_start_time:.2f} seconds")
442 + return filepath
443 +
444 +def main():
445 + parser = argparse.ArgumentParser(description="Generate music from a text prompt")
446 + parser.add_argument("prompt", nargs="?", type=str, help="Text prompt describing the desired music")
447 + parser.add_argument("--seed", type=int, default=None, help="Random seed for reproducibility")
448 + parser.add_argument("--output-dir", type=str, default=DEFAULT_OUTPUT_DIR, help="Directory to save generated music")
449 + parser.add_argument("--duration", type=int, default=None, help="Duration of music in seconds")
450 + args = parser.parse_args()
451 + if not args.prompt:
452 + print("❌ No prompt provided. Please specify a prompt.")
453 + print(f"Example: python {os.path.basename(__file__)} 'An upbeat electronic track with a catchy melody'")
454 + return 1
455 + filepath = generate_music(
456 + args.prompt,
457 + args.output_dir,
458 + duration=args.duration,
459 + seed=args.seed
460 + )
461 + if filepath:
462 + print(f"✨ Music generation completed successfully!")
463 + print_versions()
464 + try:
465 + import torch
466 + if not torch.cuda.is_available() and check_cuda():
467 + print("\n🔄 NOTE: This music was generated on CPU, but an NVIDIA GPU was detected on the host.")
468 + print("🔄 If you intended to use GPU, please check the PyTorch and CUDA driver setup within the virtual environment.")
469 + print(f"🔄 The virtual environment is located at: {VENV_DIR}")
470 + except ImportError:
471 + pass
472 + except Exception as e:
473 + print(f"Note: Error during post-generation GPU check: {e}")
474 + return 0
475 + else:
476 + print("❌ Music generation failed.")
477 + try:
478 + import torch
479 + if not torch.cuda.is_available() and check_cuda():
480 + print("\nℹ️ NOTE: An NVIDIA GPU was detected on the host, but PyTorch could not use CUDA.")
481 + print(f"ℹ️ PyTorch (version {torch.__version__}) reported CUDA as unavailable in the current runtime.")
482 + print(f"ℹ️ Dependencies (including PyTorch with CUDA if hardware was detected) were installed into: {VENV_DIR}")
483 + print("ℹ️ Please ensure your NVIDIA drivers are up to date and compatible with the PyTorch CUDA version attempted.")
484 + print("ℹ️ You might need to manually re-trigger dependency installation or debug the venv if issues persist.")
485 + except ImportError:
486 + print("ℹ️ PyTorch is not importable. Dependency installation likely failed.")
487 + except Exception as e:
488 + print(f"Note: Error during failure analysis: {e}")
489 + return 1
490 +
491 +if __name__ == "__main__":
492 + if not manage_venv_and_execution():
493 + sys.exit(1)
494 + sys.exit(main())
\ No newline at end of file
instruments/default/music_generation/music_generation.sh new
+83
@@ -0,0 +1,83 @@
1 +#!/bin/bash
2 +
3 +VENV_DIR="/opt/instruments_venv"
4 +VENV_PY="$VENV_DIR/bin/python"
5 +DEFAULT_OUTPUT_DIR="/root/generated_music"
6 +PYTHON_SCRIPT="/a0/instruments/default/music_generation/music_generation.py"
7 +
8 +echo "==== Starting Music Generation Script ===="
9 +
10 +# Show GPU info if available
11 +if command -v nvidia-smi &> /dev/null; then
12 + echo "✅ NVIDIA GPU detected, displaying information:"
13 + nvidia-smi
14 + # Get CUDA version
15 + if [ -x "$(command -v nvcc)" ]; then
16 + echo "✅ NVCC (CUDA Compiler) found:"
17 + nvcc --version
18 + else
19 + echo "⚠️ NVCC not found, CUDA development tools may not be installed properly"
20 + fi
21 + # Check CUDA libraries
22 + echo "Checking CUDA libraries:"
23 + if ldconfig -p | grep -q libcuda.so; then
24 + echo "✅ CUDA libraries found in system path"
25 + ldconfig -p | grep libcuda.so
26 + else
27 + echo "⚠️ CUDA libraries not found in system path"
28 + fi
29 +else
30 + echo "⚠️ No NVIDIA GPU detected (nvidia-smi not found)."
31 +fi
32 +
33 +# Set CUDA env vars if desired
34 +export CUDA_VISIBLE_DEVICES=0
35 +export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128
36 +
37 +# If venv Python does not exist, create venv and install CPU deps
38 +if [ ! -x "$VENV_PY" ]; then
39 + echo "🛠️ venv not found, creating at $VENV_DIR and installing CPU dependencies..."
40 + python3 -m venv "$VENV_DIR"
41 + "$VENV_PY" -m pip install --upgrade pip setuptools wheel
42 + "$VENV_PY" -m pip install \
43 + torch==2.6.0 torchvision torchaudio \
44 + huggingface-hub==0.20.3 safetensors==0.4.1 accelerate==0.21.0 \
45 + transformers==4.38.2 einops==0.6.1 tqdm==4.65.0 librosa==0.10.0.post2 scipy==1.12.0 numpy==1.24.3
46 + echo "✅ venv created and CPU dependencies installed."
47 +fi
48 +
49 +# Ensure the Python runtime can find CUDA
50 +echo "====== CUDA Environment Variables ======"
51 +echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES"
52 +echo "LD_LIBRARY_PATH: $LD_LIBRARY_PATH"
53 +
54 +echo "====== Running Music Generation ======"
55 +
56 +# Start the Python process in the background
57 +"$VENV_PY" "$PYTHON_SCRIPT" "$@" &
58 +PY_PID=$!
59 +
60 +# Heartbeat loop
61 +while kill -0 $PY_PID 2>/dev/null; do
62 + sleep 9
63 + if kill -0 $PY_PID 2>/dev/null; then
64 + echo "⏳ Process still running. Monitor terminal for output."
65 + fi
66 + # Removed repeated latest file printout from loop
67 + # (It will be printed at the end only)
68 +done
69 +
70 +wait $PY_PID
71 +status=$?
72 +
73 +if [ $status -eq 0 ]; then
74 + echo "✅ Music generation completed successfully!"
75 + LATEST_FILE=$(ls -t "$DEFAULT_OUTPUT_DIR"/*.wav 2>/dev/null | head -n 1)
76 + if [ -n "$LATEST_FILE" ]; then
77 + echo "📁 Latest generated file: $LATEST_FILE"
78 + echo "🔊 You can play this file with a media player"
79 + fi
80 +else
81 + echo "❌ Music generation failed with error code $status"
82 + exit 1
83 +fi
\ No newline at end of file