Add files via upload
Added custom nodes
This commit is contained in:
Creepybits
@@ -0,0 +1,53 @@
|
||||
import os
|
||||
import sys
|
||||
import comfy.sd
|
||||
import comfy.utils
|
||||
import re
|
||||
|
||||
class AudioKeywordExtractor:
|
||||
|
||||
def __init__(self):
|
||||
# Get the directory of the current script (SystemPromp.py)
|
||||
script_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
# Construct the full path to the text file: up one level, then into "assets/prompts"
|
||||
assets_dir = os.path.dirname(script_dir) # go up to assets
|
||||
filepath = os.path.join(assets_dir, "prompts", "audio_keywords.txt")
|
||||
|
||||
try:
|
||||
with open(filepath, "r", encoding="utf-8") as f:
|
||||
self.fixed_text = f.read()
|
||||
except FileNotFoundError:
|
||||
print(f"Error: system_prompt.txt not found at {filepath}")
|
||||
self.fixed_text = "ERROR: Prompt file not found." # Provide a fallback
|
||||
except Exception as e:
|
||||
print(f"Error reading system_prompt.txt: {e}")
|
||||
self.fixed_text = "ERROR: Could not read prompt file." # Provide a fallback
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"text_2": ("STRING", {"multiline": True}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("text",)
|
||||
|
||||
FUNCTION = "concat_texts"
|
||||
|
||||
CATEGORY = "Creepybits/Audio"
|
||||
|
||||
def concat_texts(self, text_2):
|
||||
combined_text = self.fixed_text + text_2
|
||||
return (combined_text,)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = { # <---Outdent these lines
|
||||
"AudioKeywordExtractor": AudioKeywordExtractor,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"AudioKeywordExtractor": "Audio To Image Draft (Creepybits)",
|
||||
}
|
||||
@@ -0,0 +1,385 @@
|
||||
# --- START OF FILE GeminiAudioAnalyzer.py ---
|
||||
|
||||
# Gemini_Flash_Node.py
|
||||
import os
|
||||
import json
|
||||
import base64
|
||||
import requests
|
||||
import google.generativeai as genai
|
||||
from io import BytesIO
|
||||
from PIL import Image
|
||||
import torch
|
||||
import torchaudio
|
||||
import numpy as np
|
||||
|
||||
# Define the default path for the Gemini API key file
|
||||
# Use a raw string (r"...") or double backslashes ("\\") for Windows paths
|
||||
DEFAULT_API_KEY_PATH = r"C:\AI\Comfy\ComfyUI\custom_nodes\Creepy_nodes\assets\scripts\gemini_api_key.txt"
|
||||
|
||||
# Define available safety thresholds for the dropdown
|
||||
SAFETY_THRESHOLDS = ["Block None", "Block Low", "Block Medium", "Block High"]
|
||||
|
||||
# Map dropdown values to API threshold constants (as strings for direct payload)
|
||||
# The google-generativeai library often handles the mapping of these strings internally
|
||||
SAFETY_THRESHOLD_MAP = {
|
||||
"Block None": "BLOCK_NONE",
|
||||
"Block Low": "BLOCK_LOW_AND_ABOVE",
|
||||
"Block Medium": "BLOCK_MEDIUM_AND_ABOVE",
|
||||
"Block High": "BLOCK_HIGH_AND_ABOVE",
|
||||
}
|
||||
|
||||
# Define standard safety categories to apply the threshold to
|
||||
# Using the string names as expected by the API payload / genai library
|
||||
SAFETY_CATEGORIES = [
|
||||
"HARM_CATEGORY_HARASSMENT",
|
||||
"HARM_CATEGORY_HATE_SPEECH",
|
||||
"HARM_CATEGORY_SEXUALLY_EXPLICIT",
|
||||
"HARM_CATEGORY_DANGEROUS_CONTENT",
|
||||
# Note: Some models might have additional categories like MEDICAL, DRUG, etc.
|
||||
# Check the specific model documentation if needed.
|
||||
]
|
||||
|
||||
|
||||
class GeminiAudioAnalyzer:
|
||||
"""
|
||||
A custom node for ComfyUI that uses the Google Gemini API for text generation
|
||||
with audio analysis capabilities.
|
||||
Loads API key from input field or default file.
|
||||
"""
|
||||
|
||||
CATEGORY = "Creepybits/Audio" # Category in ComfyUI interface
|
||||
RETURN_TYPES = ("STRING",) # Output data type (a single string)
|
||||
RETURN_NAMES = ("generated_content",) # Output name displayed in ComfyUI
|
||||
FUNCTION = "generate_content" # Specifies the method to run when the node is executed
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"prompt": ("STRING", {"default": "Analyze the situation in details.", "multiline": True}),
|
||||
"input_type": (["text", "audio"], {"default": "text"}), # Added text option
|
||||
"model_version": (["gemini-2.5-flash-preview-04-17", "gemini-2.5-pro-exp-03-25", "gemini-2.0-flash"], {"default": "gemini-2.0-flash"}), # Added new models
|
||||
"operation_mode": (["analysis", "generate_images"], {"default": "analysis"}), # operation_mode - seems unused in generate_content but kept as in original
|
||||
},
|
||||
"optional": {
|
||||
"Additional_Context": ("STRING", {"default": "", "multiline": True}), # Renamed from context to Additional_Context for clarity
|
||||
"audio": ("AUDIO", ), # Added audio input type
|
||||
"api_key": ("STRING", {"default": "", "multiline": False}), # API Key input field, defaults to empty string
|
||||
"max_output_tokens": ("INT", {"default": 2048, "min": 1, "max": 8192}),
|
||||
"temperature": ("FLOAT", {"default": 0.4, "min": 0.0, "max": 1.0, "step": 0.1}),
|
||||
"safety_threshold": (SAFETY_THRESHOLDS, {"default": "Block None"}),
|
||||
}
|
||||
}
|
||||
|
||||
def __init__(self):
|
||||
# Note: Chat history is not currently used in generate_content (single turn)
|
||||
self.chat_history = [] # Placeholder for chat history if multi-turn was implemented
|
||||
|
||||
def load_api_key(self, api_key=""):
|
||||
"""
|
||||
Loads the API key from the provided string parameter or the default file.
|
||||
Prioritizes the provided string.
|
||||
"""
|
||||
# Check if the api_key parameter (from the node input) is provided and not empty
|
||||
if api_key and api_key.strip(): # Added .strip() here for robustness
|
||||
# print(f"[AudioAnalyzer load_api_key] Using API key from node input.") # Debug print removed
|
||||
return api_key.strip() # Use the provided key if available, ensuring it's stripped
|
||||
|
||||
# Attempt to read from the default file path if input was empty or whitespace
|
||||
# print(f"[AudioAnalyzer load_api_key] Input empty, attempting file load from: {DEFAULT_API_KEY_PATH}") # Debug print removed
|
||||
try:
|
||||
# Attempt to read from the default file path
|
||||
with open(DEFAULT_API_KEY_PATH, "r") as f:
|
||||
file_content = f.read()
|
||||
stripped_key = file_content.strip()
|
||||
|
||||
# Check if the key read from the file is not empty after stripping
|
||||
if stripped_key:
|
||||
# print(f"[AudioAnalyzer load_api_key] File key loaded successfully.") # Debug print removed
|
||||
return stripped_key
|
||||
else:
|
||||
print(f"[AudioAnalyzer load_api_key] Warning: API key file at '{DEFAULT_API_KEY_PATH}' is empty or contains only whitespace.")
|
||||
return None # Explicitly return None if file is empty
|
||||
|
||||
except FileNotFoundError:
|
||||
print(f"[AudioAnalyzer load_api_key] Warning: API key file not found at {DEFAULT_API_KEY_PATH}. Tried to load because node input was empty.")
|
||||
return None # Return None if the file doesn't exist
|
||||
except Exception as e:
|
||||
print(f"[AudioAnalyzer load_api_key] Warning: Error reading API key file: {e}")
|
||||
return None # Return None for other file reading errors
|
||||
|
||||
|
||||
def prepare_content(self, prompt, input_type, Additional_Context=None, audio=None):
|
||||
"""Prepares the content for the Gemini API based on the input type."""
|
||||
# Combine prompt and additional context for text input
|
||||
text_content = prompt if not Additional_Context else f"{prompt}\n\n{Additional_Context.strip()}" # Added strip and separation
|
||||
|
||||
if input_type == "text":
|
||||
# For text-only input, the contents list contains a single text part
|
||||
# The genai library handles the structure if you just pass the string
|
||||
# return [{"text": text_content}] # Old way, genai library is more flexible
|
||||
return [text_content] # Pass the string directly
|
||||
|
||||
elif input_type == "audio" and audio is not None:
|
||||
# Ensure text content is included along with audio for multimodal models
|
||||
content_parts = []
|
||||
if text_content.strip(): # Only add text part if there's content
|
||||
content_parts.append({"text": text_content.strip()})
|
||||
|
||||
# Audio processing (assuming 'audio' is a dictionary with 'waveform' and 'sample_rate' from another ComfyUI node)
|
||||
try:
|
||||
waveform = audio["waveform"]
|
||||
sample_rate = audio["sample_rate"]
|
||||
|
||||
# Ensure waveform is 2D [channels, samples] for torchaudio.save
|
||||
if waveform.dim() == 3: # Assuming [batch, channels, samples] from ComfyUI AUDIO type
|
||||
waveform = waveform.squeeze(0) # Remove batch dimension
|
||||
elif waveform.dim() == 1: # Assuming [samples]
|
||||
waveform = waveform.unsqueeze(0) # Add channel dimension
|
||||
|
||||
# If multiple channels, average them (simple mono conversion)
|
||||
if waveform.shape[0] > 1:
|
||||
print(f"[AudioAnalyzer prepare_content] Warning: Input audio has {waveform.shape[0]} channels. Converting to mono by averaging.")
|
||||
waveform = torch.mean(waveform, dim=0, keepdim=True)
|
||||
|
||||
# Resample to 16kHz if needed (Gemini expects 16kHz for audio)
|
||||
if sample_rate != 16000:
|
||||
print(f"[AudioAnalyzer prepare_content] Resampling audio from {sample_rate}Hz to 16000Hz.")
|
||||
waveform = torchaudio.functional.resample(waveform, sample_rate, 16000)
|
||||
|
||||
# Save to a BytesIO buffer in WAV format
|
||||
buffer = BytesIO()
|
||||
torchaudio.save(buffer, waveform, 16000, format="WAV")
|
||||
audio_bytes = buffer.getvalue()
|
||||
|
||||
# Add the audio part as inline data
|
||||
content_parts.append({
|
||||
"inline_data": {
|
||||
"mime_type": "audio/wav",
|
||||
"data": base64.b64encode(audio_bytes).decode('utf-8') # Base64 encode the audio bytes
|
||||
}
|
||||
})
|
||||
|
||||
except KeyError as e:
|
||||
print(f"[AudioAnalyzer prepare_content] Error: Missing expected key in audio input: {e}")
|
||||
raise ValueError(f"Invalid audio input format: Missing key {e}") from e
|
||||
except Exception as e:
|
||||
print(f"[AudioAnalyzer prepare_content] Error processing audio: {e}")
|
||||
raise ValueError(f"Failed to process audio input: {e}") from e
|
||||
|
||||
# Return the list of parts (text + audio)
|
||||
return content_parts
|
||||
|
||||
else:
|
||||
# Handle cases where audio is selected but not provided, or input_type is invalid
|
||||
if input_type == "audio" and audio is None:
|
||||
raise ValueError("Input type set to 'audio', but no audio input was provided.")
|
||||
else:
|
||||
raise ValueError(f"Invalid input_type: '{input_type}'. Must be 'text' or 'audio'.")
|
||||
|
||||
|
||||
def generate_content(self, prompt, input_type, model_version="gemini-2.0-flash",
|
||||
operation_mode="analysis", Additional_Context=None, # operation_mode is unused in this function
|
||||
audio=None, api_key="", max_output_tokens=8192, # api_key parameter receives input from node field
|
||||
safety_threshold="Block None", temperature=0.4):
|
||||
"""Generate content using Gemini model with various input types (text, audio)."""
|
||||
|
||||
# --- API Key Loading ---
|
||||
# This calls load_api_key, which handles checking input field vs. file
|
||||
resolved_api_key = self.load_api_key(api_key)
|
||||
|
||||
# Check if a valid API key was successfully obtained
|
||||
if not resolved_api_key or not resolved_api_key.strip(): # Added .strip() check here too for robustness
|
||||
# print("[AudioAnalyzer generate_content] API key not resolved.") # Debug print removed
|
||||
# Return a tuple with an error message for ComfyUI string output
|
||||
return ("Error: Gemini API key not found or loaded for Audio Analyzer. Please provide key in node input, set GOOGLE_API_KEY env var, or ensure valid key in default file. Check console output for details.",)
|
||||
|
||||
# --- Configure Gemini Library ---
|
||||
try:
|
||||
# Configure the genai library with the resolved API key
|
||||
genai.configure(api_key=resolved_api_key)
|
||||
print("[AudioAnalyzer] Gemini API Key configured successfully.") # Keep this - useful confirmation
|
||||
|
||||
except Exception as e:
|
||||
# Catch potential errors during configuration (e.g., invalid key format)
|
||||
error_message_detail = str(e)
|
||||
if hasattr(e, 'response') and e.response:
|
||||
if hasattr(e.response, 'text'):
|
||||
try:
|
||||
error_json = json.loads(e.response.text)
|
||||
if 'error' in error_json and 'message' in error_json['error']:
|
||||
api_error_message = error_json['error']['message']
|
||||
error_message_detail = f"API Error during configuration: {api_error_message}"
|
||||
# Add checks for status_code/reason if needed, similar to main exception
|
||||
except (json.JSONDecodeError, AttributeError):
|
||||
pass # Keep original error message if JSON parsing fails
|
||||
|
||||
print(f"[AudioAnalyzer] Error configuring Gemini API: {error_message_detail}")
|
||||
return (f"Error configuring Gemini API: {error_message_detail}",)
|
||||
|
||||
|
||||
# --- Configure Safety Settings ---
|
||||
safety_settings = [] # Initialize safety settings list
|
||||
api_threshold = SAFETY_THRESHOLD_MAP.get(safety_threshold, "BLOCK_NONE") # Get API threshold string
|
||||
|
||||
if api_threshold != "BLOCK_NONE":
|
||||
# Apply the selected threshold to standard safety categories
|
||||
for category_name in SAFETY_CATEGORIES:
|
||||
safety_settings.append({
|
||||
"category": category_name,
|
||||
"threshold": api_threshold
|
||||
})
|
||||
print(f"[AudioAnalyzer] Applying safety threshold '{api_threshold}' to categories: {SAFETY_CATEGORIES}")
|
||||
else:
|
||||
print("[AudioAnalyzer] Safety threshold set to 'Block None'. No safety settings applied.")
|
||||
|
||||
# --- Configure Generation Parameters ---
|
||||
generation_config=genai.types.GenerationConfig(
|
||||
max_output_tokens=max_output_tokens,
|
||||
temperature=temperature,
|
||||
# Add other parameters if needed, e.g., top_p, top_k
|
||||
# top_p=...
|
||||
# top_k=...
|
||||
)
|
||||
# print(f"[AudioAnalyzer] Generation config: {generation_config}") # Optional: print full config
|
||||
|
||||
|
||||
# --- Get the Generative Model Instance ---
|
||||
try:
|
||||
# Use genai.GenerativeModel to get the specific model
|
||||
# Pass safety_settings and generation_config here
|
||||
model = genai.GenerativeModel(
|
||||
model_name=model_version,
|
||||
safety_settings=safety_settings,
|
||||
generation_config=generation_config
|
||||
)
|
||||
print(f"[AudioAnalyzer] Model instance created for '{model_version}'.") # Keep this - useful confirmation
|
||||
|
||||
except Exception as e:
|
||||
print(f"[AudioAnalyzer] Error getting model '{model_version}': {e}")
|
||||
# Check if the error is due to an invalid model name
|
||||
if "Invalid model" in str(e):
|
||||
return (f"Error: Invalid model name '{model_version}'. Check model availability.",)
|
||||
# Fallback error message
|
||||
return (f"Error: Could not get Gemini model '{model_version}'. Check model name, API status, or internet connection: {e}",)
|
||||
|
||||
|
||||
# --- Prepare Content (Text + Optional Audio) ---
|
||||
try:
|
||||
content_parts = self.prepare_content(prompt, input_type, Additional_Context, audio)
|
||||
# print(f"[AudioAnalyzer] Content parts prepared: {content_parts}") # Optional: print content parts
|
||||
except ValueError as e: # Catch specific errors from prepare_content
|
||||
print(f"[AudioAnalyzer] Error during content preparation: {e}")
|
||||
return (f"Error preparing content: {e}",)
|
||||
except Exception as e: # Catch any other unexpected errors during preparation
|
||||
print(f"[AudioAnalyzer] Unexpected error during content preparation: {e}")
|
||||
return (f"Unexpected error preparing content: {e}",)
|
||||
|
||||
|
||||
# --- Send API Request ---
|
||||
try:
|
||||
# The genai library handles the request format and sending
|
||||
print("[AudioAnalyzer] Sending generate_content request...") # Keep this - useful confirmation
|
||||
response = model.generate_content(
|
||||
contents=content_parts,
|
||||
# Safety settings and generation config are typically passed when creating the model instance
|
||||
# Passing them again here depends on the library version/design, but often not needed if done in init.
|
||||
# If your previous version worked with them here, you might need to add them back.
|
||||
# generation_config=generation_config # Example if needed
|
||||
# safety_settings=safety_settings, # Example if needed
|
||||
)
|
||||
print("[AudioAnalyzer] API request finished.") # Keep this - useful confirmation
|
||||
|
||||
|
||||
# --- Process Response ---
|
||||
generated_content = ""
|
||||
# Check if the response object has a 'text' attribute (common for success)
|
||||
if hasattr(response, 'text'):
|
||||
generated_content = response.text
|
||||
# Note: Could add more sophisticated checks here if response.text is often empty but candidates exist
|
||||
|
||||
# If no text was directly generated, check for block reasons or empty candidates
|
||||
if not generated_content:
|
||||
print("[AudioAnalyzer] Generated text is empty or extraction failed.") # Keep this warning
|
||||
|
||||
# Check prompt feedback in the response object if available (e.g., safety blocks)
|
||||
if hasattr(response, 'prompt_feedback') and response.prompt_feedback:
|
||||
print("[AudioAnalyzer] Prompt Feedback received:") # Keep this
|
||||
feedback_info = {}
|
||||
# Check if block_reason exists and is not UNASSIGNED
|
||||
if hasattr(response.prompt_feedback, 'block_reason') and response.prompt_feedback.block_reason and hasattr(response.prompt_feedback.block_reason, 'name') and response.prompt_feedback.block_reason.name != 'UNASSIGNED':
|
||||
print(f" Block Reason: {response.prompt_feedback.block_reason.name}") # Keep this
|
||||
feedback_info['block_reason'] = response.prompt_feedback.block_reason.name
|
||||
|
||||
if hasattr(response.prompt_feedback, 'safety_ratings') and response.prompt_feedback.safety_ratings:
|
||||
print(" Safety Ratings:") # Keep this
|
||||
feedback_info['safety_ratings'] = []
|
||||
for rating in response.prompt_feedback.safety_ratings:
|
||||
# Only print/report safety ratings that are not UNSPECIFIED probability
|
||||
if hasattr(rating, 'probability') and hasattr(rating.probability, 'name') and rating.probability.name != 'UNSPECIFIED':
|
||||
rating_info = {"category": rating.category.name, "probability": rating.probability.name}
|
||||
print(f" Category: {rating_info['category']}, Probability: {rating_info['probability']}") # Keep this
|
||||
feedback_info['safety_ratings'].append(rating_info)
|
||||
|
||||
# Provide a user-friendly error message based on feedback
|
||||
if 'block_reason' in feedback_info:
|
||||
return (f"Error: Content generation blocked due to {feedback_info['block_reason']}. Check console for safety feedback.",)
|
||||
elif 'safety_ratings' in feedback_info and feedback_info['safety_ratings']:
|
||||
# If no block reason, but safety ratings are high, it still might have influenced output
|
||||
return ("Error: No text generated. High safety ratings. Check console for details.",)
|
||||
# Fallback if prompt_feedback exists but doesn't have block reason or significant safety issues
|
||||
return ("Error: No text generated. API returned empty content/candidates. Check Prompt Feedback in console.",)
|
||||
else:
|
||||
# If no prompt feedback object, or it's empty, or doesn't have expected info
|
||||
# Check if candidates list was explicitly empty (often means blocked)
|
||||
if hasattr(response, 'candidates') and not response.candidates:
|
||||
return ("Error: API returned no candidates (likely blocked). No prompt feedback details available.",)
|
||||
# If candidates exist but text wasn't extracted, or other unknown issue
|
||||
return ("Error: No text generated from the API and no specific feedback provided. Check console.",)
|
||||
|
||||
|
||||
# print(f"[AudioAnalyzer] Generated content length: {len(generated_content)}") # Optional: print length
|
||||
# ComfyUI nodes returning STRING should return a tuple containing the string
|
||||
return (generated_content,)
|
||||
|
||||
# --- Error Handling (for exceptions during API call) ---
|
||||
except Exception as e:
|
||||
# Catch any exception during the API call itself
|
||||
print(f"[AudioAnalyzer] An error occurred during Gemini API call: {e}") # Keep this basic error print
|
||||
|
||||
# Attempt to get more specific error details from the exception if it's an API error
|
||||
error_message_detail = str(e) # Default error message is the exception string
|
||||
if hasattr(e, 'response') and e.response:
|
||||
if hasattr(e.response, 'text'):
|
||||
try:
|
||||
error_json = json.loads(e.response.text)
|
||||
if 'error' in error_json and 'message' in error_json['error']:
|
||||
api_error_message = error_json['error']['message']
|
||||
print(f"[AudioAnalyzer] API Error Message from response: {api_error_message}") # Keep this
|
||||
error_message_detail = f"API Error: {api_error_message}"
|
||||
elif hasattr(e.response, 'status_code') and hasattr(e.response, 'reason'):
|
||||
# Catch standard HTTP error details if not JSON
|
||||
error_message_detail = f"HTTP Error: {e.response.status_code} {e.response.reason}"
|
||||
except (json.JSONDecodeError, AttributeError):
|
||||
# If response text is not JSON or AttributeError accessing parts
|
||||
if hasattr(e.response, 'status_code') and hasattr(e.response, 'reason'):
|
||||
error_message_detail = f"HTTP Error: {e.response.status_code} {e.response.reason}"
|
||||
pass # Use the original exception string if cannot extract API specific
|
||||
|
||||
elif hasattr(e.response, 'status_code') and hasattr(e.response, 'reason'):
|
||||
error_message_detail = f"HTTP Error: {e.response.status_code} {e.response.reason}"
|
||||
|
||||
# Return a tuple with the error message for ComfyUI string output
|
||||
return (f"Error: An error occurred during Gemini API call: {error_message_detail}",)
|
||||
|
||||
|
||||
# --- ComfyUI Node Mappings ---
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"GeminiAudioAnalyzer": GeminiAudioAnalyzer,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"GeminiAudioAnalyzer": "Gemini Audio Analyzer (Creepybits)",
|
||||
}
|
||||
|
||||
# --- END OF FILE GeminiAudioAnalyzer.py ---
|
||||
@@ -0,0 +1,143 @@
|
||||
# --- START OF FILE RandomAudioSegment.py ---
|
||||
|
||||
import torch
|
||||
import torchaudio
|
||||
import random
|
||||
import math # Import math for isnan
|
||||
|
||||
class RandomAudioSegment:
|
||||
CATEGORY = "Creepybits/Audio" # Or your preferred category
|
||||
RETURN_TYPES = ("AUDIO",)
|
||||
RETURN_NAMES = ("audio",)
|
||||
FUNCTION = "get_random_segment"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"audio": ("AUDIO",),
|
||||
"segment_length": ("FLOAT", {
|
||||
"default": 10.0,
|
||||
"min": 0.1, # Minimum segment length
|
||||
"max": 600.0, # Maximum segment length (e.g., 10 minutes)
|
||||
"step": 0.1,
|
||||
"display": "number" # Explicitly hint as a number input
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
# Add the new input for fixed start time
|
||||
# Use -1.0 as a sentinel value to indicate "use random start time"
|
||||
"start_time": ("FLOAT", {
|
||||
"default": -1.0, # Default to -1.0 to trigger random selection
|
||||
"min": -1.0, # Allow -1.0 (or any negative) for random, 0.0 or positive for fixed
|
||||
"step": 0.1,
|
||||
"display": "number"
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
# Update the function signature to accept the new parameter
|
||||
def get_random_segment(self, audio, segment_length=10.0, start_time=-1.0): # Add start_time with default -1.0
|
||||
waveform = audio["waveform"]
|
||||
sample_rate = audio["sample_rate"]
|
||||
|
||||
# Ensure segment_length is positive
|
||||
if segment_length <= 0:
|
||||
print(f"[RandomAudioSegment] Warning: segment_length must be positive ({segment_length} provided). Using default 10.0.")
|
||||
segment_length = 10.0
|
||||
# Handle potential NaN from UI
|
||||
if math.isnan(segment_length):
|
||||
print(f"[RandomAudioSegment] Warning: segment_length is NaN. Using default 10.0.")
|
||||
segment_length = 10.0
|
||||
|
||||
|
||||
# Calculate total duration in seconds
|
||||
total_duration = waveform.shape[-1] / sample_rate
|
||||
|
||||
# Determine the actual segment length to extract (capped by total duration)
|
||||
actual_segment_length = min(segment_length, total_duration)
|
||||
|
||||
# Calculate maximum possible start time for this actual segment length
|
||||
max_possible_start_time = total_duration - actual_segment_length
|
||||
|
||||
selected_start_time = 0.0 # Initialize selected start time in seconds
|
||||
|
||||
# Check if a fixed start time is provided (start_time >= 0)
|
||||
if start_time >= 0.0 and not math.isnan(start_time): # Also check for NaN
|
||||
print(f"[RandomAudioSegment] Fixed start time provided: {start_time:.2f}s")
|
||||
selected_start_time = start_time
|
||||
|
||||
# Clamp the requested start time to the valid range if necessary
|
||||
if max_possible_start_time > 0 and selected_start_time > max_possible_start_time:
|
||||
print(f"[RandomAudioSegment] Warning: Requested start time {start_time:.2f}s exceeds maximum possible start time {max_possible_start_time:.2f}s for segment length {actual_segment_length:.2f}s. Clamping start time.")
|
||||
selected_start_time = max_possible_start_time
|
||||
elif max_possible_start_time <= 0 and selected_start_time > 0:
|
||||
# Handle case where audio is shorter than segment, but user gave positive start time
|
||||
print(f"[RandomAudioSegment] Warning: Audio ({total_duration:.2f}s) is shorter than or equal to segment length ({actual_segment_length:.2f}s). Fixed start time > 0 is invalid. Starting at 0.")
|
||||
selected_start_time = 0.0
|
||||
|
||||
|
||||
else:
|
||||
# If start_time is negative or NaN, use random start time logic
|
||||
print(f"[RandomAudioSegment] Using random start time (start_time input was {start_time}).")
|
||||
if max_possible_start_time <= 0:
|
||||
# If audio is shorter than or equal to the desired segment length, start at 0
|
||||
selected_start_time = 0.0
|
||||
# The actual segment length will be the total duration in this case
|
||||
actual_segment_length = total_duration
|
||||
print(f"[RandomAudioSegment] Audio too short for requested segment length. Using full audio from start: 0.0s, length: {actual_segment_length:.2f}s.")
|
||||
else:
|
||||
# Otherwise, select a random start time within the valid range
|
||||
selected_start_time = random.uniform(0.0, max_possible_start_time)
|
||||
print(f"[RandomAudioSegment] Random start time selected: {selected_start_time:.2f}s, length: {actual_segment_length:.2f}s.")
|
||||
|
||||
# Convert start time to number of samples
|
||||
start_sample = int(selected_start_time * sample_rate)
|
||||
|
||||
# Calculate end sample based on start_sample and actual_segment_length
|
||||
num_samples = int(actual_segment_length * sample_rate)
|
||||
end_sample = start_sample + num_samples
|
||||
|
||||
# Handle potential off-by-one or floating point issues at the end
|
||||
# Ensure end_sample does not exceed the total number of samples in the original waveform
|
||||
end_sample = min(end_sample, waveform.shape[-1])
|
||||
|
||||
# Recalculate actual number of samples extracted based on the clamped end_sample
|
||||
actual_num_samples_extracted = end_sample - start_sample
|
||||
|
||||
# Basic check: if actual_num_samples_extracted ended up 0 or negative, return an empty segment
|
||||
if actual_num_samples_extracted <= 0:
|
||||
print(f"[RandomAudioSegment] Warning: Calculated segment resulted in zero or less samples. Returning empty audio.")
|
||||
# Return an empty tensor with the correct shape (channels, 0) or (batch, channels, 0)
|
||||
# Need to preserve leading dimensions (batch, channels)
|
||||
segment = torch.empty(*waveform.shape[:-1], 0, dtype=waveform.dtype, device=waveform.device)
|
||||
extracted_duration = 0.0 # Duration is zero for empty segment
|
||||
else:
|
||||
# Extract segment using calculated samples
|
||||
# Assuming waveform is [batch, channels, samples] or [channels, samples]
|
||||
# Use slicing on the last dimension (samples)
|
||||
segment = waveform[..., start_sample:end_sample]
|
||||
extracted_duration = actual_num_samples_extracted / sample_rate # Calculate actual duration extracted
|
||||
|
||||
|
||||
print(f"[RandomAudioSegment] Extracted segment: Start Sample={start_sample}, End Sample={end_sample}, Actual Samples={actual_num_samples_extracted}, Duration={extracted_duration:.2f}s")
|
||||
|
||||
|
||||
# Create new audio dictionary
|
||||
new_audio = {
|
||||
"waveform": segment,
|
||||
"sample_rate": sample_rate,
|
||||
}
|
||||
|
||||
return (new_audio,) # Return as a tuple
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"RandomAudioSegment": RandomAudioSegment,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"RandomAudioSegment": "Random/Fixed Audio Picker (Creepybits)", # Updated display name
|
||||
}
|
||||
|
||||
# --- END OF FILE RandomAudioSegment.py ---
|
||||
Reference in New Issue
Block a user