Spaces:
Running on Zero
Running on Zero
| import os | |
| import random | |
| import requests | |
| import hashlib | |
| import re | |
| from typing import Sequence, Mapping, Any, Union, Set | |
| from pathlib import Path | |
| import shutil | |
| import gradio as gr | |
| from huggingface_hub import hf_hub_download, constants as hf_constants | |
| import torch | |
| import numpy as np | |
| from PIL import Image, ImageChops | |
| import yaml | |
| from core.settings import * | |
| MODELS_ROOT_DIR = "ComfyUI/models" | |
| class UniqueKeyLoader(yaml.SafeLoader): | |
| """ | |
| A custom YAML loader that handles duplicate keys by grouping their values into a list. | |
| """ | |
| def construct_mapping(self, node, deep=False): | |
| mapping = [] | |
| for key_node, value_node in node.value: | |
| key = self.construct_object(key_node, deep=deep) | |
| value = self.construct_object(value_node, deep=deep) | |
| mapping.append((key, value)) | |
| result = {} | |
| for k, v in mapping: | |
| if k in result: | |
| if isinstance(result[k], list): | |
| result[k].append(v) | |
| else: | |
| result[k] = [result[k], v] | |
| else: | |
| result[k] = v | |
| return result | |
| UniqueKeyLoader.add_constructor(yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, UniqueKeyLoader.construct_mapping) | |
| def save_uploaded_file_with_hash(file_obj: gr.File, target_dir: str) -> str: | |
| if not file_obj: | |
| return "" | |
| temp_path = file_obj.name | |
| sha256 = hashlib.sha256() | |
| with open(temp_path, 'rb') as f: | |
| for block in iter(lambda: f.read(65536), b''): | |
| sha256.update(block) | |
| file_hash = sha256.hexdigest() | |
| _, extension = os.path.splitext(temp_path) | |
| hashed_filename = f"{file_hash}{extension.lower()}" | |
| dest_path = os.path.join(target_dir, hashed_filename) | |
| os.makedirs(target_dir, exist_ok=True) | |
| if not os.path.exists(dest_path): | |
| shutil.copy(temp_path, dest_path) | |
| print(f"✅ Saved uploaded file as: {dest_path}") | |
| else: | |
| print(f"ℹ️ File already exists (deduplicated): {dest_path}") | |
| return hashed_filename | |
| def bytes_to_gb(byte_size: int) -> float: | |
| if byte_size is None or byte_size == 0: | |
| return 0.0 | |
| return round(byte_size / (1024 ** 3), 2) | |
| def get_directory_size(path: str) -> int: | |
| total_size = 0 | |
| if not os.path.exists(path): | |
| return 0 | |
| try: | |
| for dirpath, _, filenames in os.walk(path): | |
| for f in filenames: | |
| fp = os.path.join(dirpath, f) | |
| if os.path.isfile(fp) and not os.path.islink(fp): | |
| total_size += os.path.getsize(fp) | |
| except OSError as e: | |
| print(f"Warning: Could not access {path} to calculate size: {e}") | |
| return total_size | |
| def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any: | |
| try: | |
| return obj[index] | |
| except (KeyError, IndexError): | |
| try: | |
| return obj["result"][index] | |
| except (KeyError, IndexError): | |
| return None | |
| def sanitize_prompt(prompt: str) -> str: | |
| if not isinstance(prompt, str): | |
| return "" | |
| return "".join(char for char in prompt if char.isprintable() or char in ('\n', '\t')) | |
| def sanitize_id(input_id: str) -> str: | |
| if not isinstance(input_id, str): | |
| return "" | |
| input_id = input_id.strip() | |
| if "civitai" in input_id.lower(): | |
| version_match = re.search(r'modelVersionId=(\d+)', input_id) | |
| if version_match: | |
| return version_match.group(1) | |
| model_match = re.search(r'/models/(\d+)', input_id) | |
| if model_match: | |
| return model_match.group(1) | |
| return re.sub(r'[^0-9]', '', input_id) | |
| def sanitize_url(url: str) -> str: | |
| if not isinstance(url, str): | |
| raise ValueError("URL must be a string.") | |
| url = url.strip() | |
| if not re.match(r'^https?://[^\s/$.?#].[^\s]*$', url): | |
| raise ValueError("Invalid URL format or scheme. Only HTTP and HTTPS are allowed.") | |
| return url | |
| def sanitize_filename(filename: str) -> str: | |
| if not isinstance(filename, str): | |
| return "" | |
| sanitized = filename.replace('..', '') | |
| sanitized = re.sub(r'[^\w\.\-]', '_', sanitized) | |
| return sanitized.lstrip('/\\') | |
| def get_civitai_file_info(version_id: str) -> dict | None: | |
| api_url = f"https://civitai.com/api/v1/model-versions/{version_id}" | |
| try: | |
| response = requests.get(api_url, timeout=10) | |
| response.raise_for_status() | |
| data = response.json() | |
| model_type = data.get('model', {}).get('type') | |
| result_file = None | |
| for file_data in data.get('files', []): | |
| if file_data.get('type') == 'Model' and file_data['name'].endswith(('.safetensors', '.pt', '.bin')): | |
| result_file = file_data.copy() | |
| break | |
| if not result_file and data.get('files'): | |
| result_file = data['files'][0].copy() | |
| if result_file: | |
| result_file['model_type'] = model_type | |
| return result_file | |
| except Exception: | |
| return None | |
| def download_file(url: str, save_path: str, api_key: str = None, progress=None, desc: str = "") -> str: | |
| if os.path.exists(save_path): | |
| return f"File already exists: {os.path.basename(save_path)}" | |
| headers = {'Authorization': f'Bearer {api_key}'} if api_key and api_key.strip() else {} | |
| try: | |
| if progress: | |
| progress(0, desc=desc) | |
| response = requests.get(url, stream=True, headers=headers, timeout=15) | |
| response.raise_for_status() | |
| total_size = int(response.headers.get('content-length', 0)) | |
| with open(save_path, "wb") as f: | |
| downloaded = 0 | |
| for chunk in response.iter_content(chunk_size=8192): | |
| f.write(chunk) | |
| if progress and total_size > 0: | |
| downloaded += len(chunk) | |
| progress(downloaded / total_size, desc=desc) | |
| return f"Successfully downloaded: {os.path.basename(save_path)}" | |
| except Exception as e: | |
| if os.path.exists(save_path): | |
| os.remove(save_path) | |
| return f"Download failed for {os.path.basename(save_path)}: {e}" | |
| def get_lora_path(source: str, id_or_url: str, civitai_key: str, progress) -> tuple[str | None, str]: | |
| if not id_or_url or not id_or_url.strip(): | |
| return None, "No ID/URL provided." | |
| try: | |
| if source == "Civitai": | |
| version_id = sanitize_id(id_or_url) | |
| if not version_id: | |
| return None, "Invalid Civitai ID provided. Must be numeric." | |
| file_info = get_civitai_file_info(version_id) | |
| if file_info: | |
| model_type = file_info.get('model_type') | |
| if model_type and model_type.lower() == 'checkpoint': | |
| return None, f"Invalid Civitai model type '{model_type}' for LoRA. Checkpoint models are not allowed." | |
| filename = sanitize_filename(f"civitai_{version_id}.safetensors") | |
| local_path = os.path.join(LORA_DIR, filename) | |
| api_key_to_use = civitai_key | |
| source_name = f"Civitai ID {version_id}" | |
| elif source == "Hugging Face": | |
| parts = id_or_url.strip().split('/') | |
| if len(parts) < 3: | |
| return None, "Invalid Hugging Face path. Format: repo_owner/repo_name/filename" | |
| repo_id = f"{parts[0]}/{parts[1]}" | |
| repo_file_path = "/".join(parts[2:]) | |
| unique_name = id_or_url.strip().replace('/', '_') | |
| filename = sanitize_filename(unique_name) | |
| local_path = os.path.join(LORA_DIR, filename) | |
| source_name = f"HF {repo_file_path}" | |
| else: | |
| return None, "Invalid source." | |
| except ValueError as e: | |
| return None, f"Input validation failed: {e}" | |
| if os.path.lexists(local_path): | |
| if not os.path.exists(local_path): | |
| os.remove(local_path) | |
| else: | |
| return local_path, "File already exists." | |
| if source == "Civitai": | |
| if not file_info or not file_info.get('downloadUrl'): | |
| return None, f"Could not get download link for {source_name}." | |
| status = download_file(file_info['downloadUrl'], local_path, api_key_to_use, progress=progress, desc=f"Downloading {source_name}") | |
| return (local_path, status) if "Successfully" in status else (None, status) | |
| elif source == "Hugging Face": | |
| try: | |
| if progress and callable(progress): progress(0, desc=f"Downloading {source_name}") | |
| cached_path = hf_hub_download(repo_id=repo_id, filename=repo_file_path, token=os.environ.get("HF_TOKEN")) | |
| os.makedirs(LORA_DIR, exist_ok=True) | |
| if os.path.lexists(local_path): | |
| if not os.path.exists(local_path): | |
| try: | |
| os.remove(local_path) | |
| except OSError: | |
| pass | |
| if not os.path.exists(local_path): | |
| try: | |
| os.symlink(cached_path, local_path) | |
| except (OSError, NotImplementedError): | |
| shutil.copyfile(cached_path, local_path) | |
| if progress and callable(progress): progress(1.0, desc=f"Downloaded {source_name}") | |
| return local_path, f"Successfully downloaded: {filename}" | |
| except Exception as e: | |
| return None, f"Hugging Face download failed: {e}" | |
| def _ensure_model_downloaded(display_name: str, progress=gr.Progress()): | |
| if display_name not in ALL_MODEL_MAP: | |
| for cat_dir in CATEGORY_TO_DIR_MAP.values(): | |
| check_path = os.path.join(cat_dir, display_name) | |
| if os.path.exists(check_path): | |
| return display_name | |
| raise ValueError(f"Model '{display_name}' not found in configuration.") | |
| model_info = ALL_MODEL_MAP[display_name] | |
| repo_filename = model_info[1] | |
| base_filename = os.path.basename(repo_filename) | |
| download_info = ALL_FILE_DOWNLOAD_MAP.get(base_filename) | |
| if not download_info: | |
| raise gr.Error(f"Model '{base_filename}' not found in file_list.yaml. Cannot download.") | |
| category = download_info.get("category") | |
| dest_dir = CATEGORY_TO_DIR_MAP.get(category) | |
| if not dest_dir: | |
| raise ValueError(f"Unknown YAML category '{category}' for '{base_filename}'.") | |
| dest_path = os.path.join(dest_dir, base_filename) | |
| if os.path.lexists(dest_path): | |
| if not os.path.exists(dest_path): | |
| print(f"⚠️ Found and removed broken symlink: {dest_path}") | |
| os.remove(dest_path) | |
| else: | |
| return base_filename | |
| source = download_info.get("source") | |
| try: | |
| progress(0, desc=f"Downloading: {base_filename}") | |
| if source == "hf": | |
| repo_id = download_info.get("repo_id") | |
| hf_filename = download_info.get("repository_file_path", base_filename) | |
| if not repo_id: | |
| raise ValueError(f"repo_id is missing for HF model '{base_filename}'") | |
| cached_path = hf_hub_download(repo_id=repo_id, filename=hf_filename, token=os.environ.get("HF_TOKEN")) | |
| os.makedirs(dest_dir, exist_ok=True) | |
| os.symlink(cached_path, dest_path) | |
| print(f"✅ Symlinked '{cached_path}' to '{dest_path}'") | |
| elif source == "civitai": | |
| model_version_id = download_info.get("model_version_id") | |
| if not model_version_id: | |
| raise ValueError(f"model_version_id is missing for Civitai model '{base_filename}'") | |
| file_info = get_civitai_file_info(model_version_id) | |
| if not file_info or not file_info.get('downloadUrl'): | |
| raise ConnectionError(f"Could not get download URL for Civitai model version ID {model_version_id}") | |
| status = download_file( | |
| file_info['downloadUrl'], dest_path, api_key=os.environ.get("CIVITAI_API_KEY", ""), progress=progress, desc=f"Downloading: {base_filename}" | |
| ) | |
| if "Failed" in status: | |
| raise ConnectionError(status) | |
| else: | |
| raise NotImplementedError(f"Download source '{source}' is not implemented for '{base_filename}'") | |
| progress(1.0, desc=f"Downloaded: {base_filename}") | |
| except Exception as e: | |
| if os.path.lexists(dest_path): | |
| try: | |
| os.remove(dest_path) | |
| except OSError: pass | |
| raise gr.Error(f"Failed to download and link '{display_name}': {e}") | |
| return base_filename | |
| def ensure_file_downloaded(filename: str, progress=None): | |
| if not filename or filename == "None": | |
| return | |
| download_info = ALL_FILE_DOWNLOAD_MAP.get(filename) | |
| if not download_info: | |
| print(f"⚠️ Warning: File '{filename}' not found in configuration (file_list.yaml). Cannot download.") | |
| return | |
| category = download_info.get("category", "loras") | |
| dest_dir = CATEGORY_TO_DIR_MAP.get(category, LORA_DIR) | |
| dest_path = os.path.join(dest_dir, filename) | |
| if os.path.lexists(dest_path): | |
| if not os.path.exists(dest_path): | |
| print(f"⚠️ Found and removed broken symlink: {dest_path}") | |
| os.remove(dest_path) | |
| else: | |
| return | |
| source = download_info.get("source") | |
| try: | |
| if source == "hf": | |
| repo_id = download_info.get("repo_id") | |
| repo_filename = download_info.get("repository_file_path", filename) | |
| if not repo_id: | |
| raise ValueError("repo_id is missing for Hugging Face download.") | |
| if progress and callable(progress): | |
| progress(0, desc=f"Downloading: {filename}") | |
| cached_path = hf_hub_download(repo_id=repo_id, filename=repo_filename, token=os.environ.get("HF_TOKEN")) | |
| os.makedirs(dest_dir, exist_ok=True) | |
| os.symlink(cached_path, dest_path) | |
| print(f"✅ Symlinked '{cached_path}' to '{dest_path}'") | |
| if progress and callable(progress): | |
| progress(1.0, desc=f"Downloaded: {filename}") | |
| elif source == "civitai": | |
| model_version_id = download_info.get("model_version_id") | |
| if not model_version_id: | |
| raise ValueError("model_version_id is missing for Civitai download.") | |
| file_info = get_civitai_file_info(model_version_id) | |
| if not file_info or not file_info.get('downloadUrl'): | |
| raise ConnectionError(f"Could not get download URL for Civitai model version ID {model_version_id}") | |
| status = download_file( | |
| file_info['downloadUrl'], | |
| dest_path, | |
| api_key=os.environ.get("CIVITAI_API_KEY", ""), | |
| progress=progress, | |
| desc=f"Downloading: {filename}" | |
| ) | |
| if "Failed" in status: | |
| raise ConnectionError(status) | |
| else: | |
| raise NotImplementedError(f"Download source '{source}' is not implemented for '{filename}'.") | |
| except Exception as e: | |
| if os.path.lexists(dest_path): | |
| try: | |
| os.remove(dest_path) | |
| except OSError: | |
| pass | |
| raise gr.Error(f"Failed to download file '{filename}': {e}") | |
| def get_model_generation_defaults(model_display_name: str, model_type: str, defaults_config: dict): | |
| final_defaults = { | |
| 'steps': 25, 'cfg': 7.0, 'sampler_name': 'euler', 'scheduler': 'simple', | |
| 'positive_prompt': '', 'negative_prompt': '' | |
| } | |
| if 'Default' in defaults_config: | |
| final_defaults.update(defaults_config['Default']) | |
| model_type_key = next((key for key in defaults_config if key.lower().replace(" ", "-").replace(".", "") == model_type.lower()), None) | |
| if model_type_key: | |
| model_type_config = defaults_config[model_type_key] | |
| if '_defaults' in model_type_config: | |
| final_defaults.update(model_type_config['_defaults']) | |
| if model_display_name in model_type_config: | |
| final_defaults.update(model_type_config[model_display_name]) | |
| return final_defaults | |
| def get_filename_prefix() -> str: | |
| import time | |
| return f"H3_{int(time.time())}" | |
| def save_temp_image(img): | |
| if img is None: | |
| return None | |
| _PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) | |
| input_dir = os.path.join(_PROJECT_ROOT, "input") | |
| os.makedirs(input_dir, exist_ok=True) | |
| if isinstance(img, Image.Image): | |
| filename = f"temp_image_{random.randint(10000, 99999)}.png" | |
| filepath = os.path.join(input_dir, filename) | |
| img.save(filepath, "PNG") | |
| return os.path.basename(filepath) | |
| elif isinstance(img, str): | |
| if not img: | |
| return None | |
| if os.path.exists(img): | |
| ext = os.path.splitext(img)[1] or ".png" | |
| filename = f"temp_image_{random.randint(10000, 99999)}{ext}" | |
| save_path = os.path.join(input_dir, filename) | |
| shutil.copy(img, save_path) | |
| return os.path.basename(save_path) | |
| if os.path.exists(os.path.join(input_dir, img)): | |
| return img | |
| return os.path.basename(img) | |
| return None | |
| def save_temp_audio(audio_path): | |
| if not audio_path: | |
| return None | |
| _PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) | |
| input_dir = os.path.join(_PROJECT_ROOT, "input") | |
| os.makedirs(input_dir, exist_ok=True) | |
| if os.path.exists(audio_path): | |
| ext = os.path.splitext(audio_path)[1] or ".wav" | |
| filename = f"temp_audio_{random.randint(10000, 99999)}{ext}" | |
| save_path = os.path.join(input_dir, filename) | |
| shutil.copy(audio_path, save_path) | |
| return os.path.basename(filename) | |
| if os.path.exists(os.path.join(input_dir, audio_path)): | |
| return audio_path | |
| return os.path.basename(audio_path) | |
| def save_temp_video(video_path): | |
| if not video_path: | |
| return None | |
| _PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) | |
| input_dir = os.path.join(_PROJECT_ROOT, "input") | |
| os.makedirs(input_dir, exist_ok=True) | |
| if os.path.exists(video_path): | |
| ext = os.path.splitext(video_path)[1] or ".mp4" | |
| filename = f"temp_video_{random.randint(10000, 99999)}{ext}" | |
| save_path = os.path.join(input_dir, filename) | |
| shutil.copy(video_path, save_path) | |
| return os.path.basename(filename) | |
| if os.path.exists(os.path.join(input_dir, video_path)): | |
| return video_path | |
| return os.path.basename(video_path) | |
| def handle_seed(seed_value: int, max_val: int = 2**32 - 1) -> int: | |
| if seed_value == -1 or seed_value is None: | |
| return random.randint(0, max_val) | |
| return int(seed_value) | |
| def process_lora_inputs(ui_values: dict, prefix: str = "", progress=None) -> list: | |
| active_loras_for_gpu = [] | |
| # 1. Check direct prefix format (e.g. lora_sources_h3_fl2va) | |
| lora_sources = ui_values.get(f'lora_sources_{prefix}', []) if prefix else [] | |
| lora_ids = ui_values.get(f'lora_ids_{prefix}', []) if prefix else [] | |
| lora_scales = ui_values.get(f'lora_scales_{prefix}', []) if prefix else [] | |
| if isinstance(lora_sources, list) and isinstance(lora_ids, list): | |
| for source, val, scale in zip(lora_sources, lora_ids, lora_scales): | |
| scale_val = float(scale) if scale is not None else 1.0 | |
| if scale_val > 0 and val and str(val).strip(): | |
| lora_id = str(val).strip() | |
| lora_filename = None | |
| if source == "File": | |
| lora_filename = sanitize_filename(lora_id) | |
| local_path = os.path.join(LORA_DIR, lora_filename) | |
| if not os.path.exists(local_path): | |
| raise gr.Error(f"Uploaded LoRA file '{lora_id}' no longer exists on server. Please re-upload it.") | |
| elif source in ("Civitai", "Hugging Face"): | |
| local_path, status = get_lora_path(source, lora_id, os.environ.get("CIVITAI_API_KEY", ""), progress) | |
| if local_path: | |
| lora_filename = os.path.basename(local_path) | |
| else: | |
| raise gr.Error(f"Failed to prepare LoRA {lora_id}: {status}") | |
| if lora_filename: | |
| active_loras_for_gpu.append({ | |
| "lora_name": lora_filename, | |
| "strength_model": scale_val, | |
| "strength_clip": scale_val | |
| }) | |
| # 2. Check lora_data flat list format (e.g. [source1, id1, scale1, upload1, ...]) | |
| lora_data = ui_values.get('lora_data', []) | |
| if lora_data and not active_loras_for_gpu: | |
| sources, ids, scales, files = lora_data[0::4], lora_data[1::4], lora_data[2::4], lora_data[3::4] | |
| for source, lora_id, scale, _ in zip(sources, ids, scales, files): | |
| scale_val = float(scale) if scale is not None else 1.0 | |
| if scale_val > 0 and lora_id and str(lora_id).strip(): | |
| lora_id_str = str(lora_id).strip() | |
| lora_filename = None | |
| if source == "File": | |
| lora_filename = sanitize_filename(lora_id_str) | |
| local_path = os.path.join(LORA_DIR, lora_filename) | |
| if not os.path.exists(local_path): | |
| raise gr.Error(f"Uploaded LoRA file '{lora_id_str}' no longer exists on server. Please re-upload it.") | |
| elif source in ("Civitai", "Hugging Face"): | |
| local_path, status = get_lora_path(source, lora_id_str, os.environ.get("CIVITAI_API_KEY", ""), progress) | |
| if local_path: | |
| lora_filename = os.path.basename(local_path) | |
| else: | |
| raise gr.Error(f"Failed to prepare LoRA {lora_id_str}: {status}") | |
| if lora_filename: | |
| active_loras_for_gpu.append({ | |
| "lora_name": lora_filename, | |
| "strength_model": scale_val, | |
| "strength_clip": scale_val | |
| }) | |
| # 3. Check direct 'loras' list of dicts (from MCP or custom payload) | |
| raw_loras = ui_values.get('loras', []) | |
| if raw_loras and not active_loras_for_gpu and isinstance(raw_loras, list): | |
| for item in raw_loras: | |
| if isinstance(item, dict): | |
| if "lora_name" in item: | |
| active_loras_for_gpu.append(item) | |
| else: | |
| src = item.get("source", "Hugging Face") | |
| val = item.get("lora_value") or item.get("id_or_url") or item.get("lora_id") | |
| scale = item.get("scale", 1.0) | |
| scale_val = float(scale) if scale is not None else 1.0 | |
| if scale_val > 0 and val and str(val).strip(): | |
| lora_id_str = str(val).strip() | |
| lora_filename = None | |
| if src == "File": | |
| lora_filename = sanitize_filename(lora_id_str) | |
| local_path = os.path.join(LORA_DIR, lora_filename) | |
| if not os.path.exists(local_path): | |
| raise gr.Error(f"Uploaded LoRA file '{lora_id_str}' no longer exists on server. Please re-upload it.") | |
| elif src in ("Civitai", "Hugging Face"): | |
| local_path, status = get_lora_path(src, lora_id_str, os.environ.get("CIVITAI_API_KEY", ""), progress) | |
| if local_path: | |
| lora_filename = os.path.basename(local_path) | |
| else: | |
| raise gr.Error(f"Failed to prepare LoRA {lora_id_str}: {status}") | |
| if lora_filename: | |
| active_loras_for_gpu.append({ | |
| "lora_name": lora_filename, | |
| "strength_model": scale_val, | |
| "strength_clip": scale_val | |
| }) | |
| return active_loras_for_gpu | |
| def load_h3_controlnet_config(): | |
| _PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) | |
| _CN_MODEL_LIST_PATH = os.path.join(_PROJECT_ROOT, 'yaml', 'h3_controlnet_models.yaml') | |
| try: | |
| with open(_CN_MODEL_LIST_PATH, 'r', encoding='utf-8') as f: | |
| config = yaml.safe_load(f) | |
| return config.get("H3_ControlNet", {}) if isinstance(config, dict) else {} | |
| except Exception as e: | |
| print(f"Error loading h3_controlnet_models.yaml: {e}") | |
| return {} | |
| def get_h3_cn_defaults(arch_val="MiniMax-H3"): | |
| cn_full_config = load_h3_controlnet_config() | |
| cn_config = cn_full_config.get(arch_val, []) | |
| if not cn_config and cn_full_config: | |
| cn_config = next(iter(cn_full_config.values()), []) | |
| if not cn_config: | |
| return ["Canny", "Depth", "HED", "MLSD", "Pose"], "Canny", ["alibaba-pai/MiniMax-H3-Fun-Controlnet-Union"], "alibaba-pai/MiniMax-H3-Fun-Controlnet-Union", "minimax_h3_fun_controlnet_union_pruned_int8_convrot.safetensors" | |
| all_types = [] | |
| for model in cn_config: | |
| for t in model.get("Type", []): | |
| if t not in all_types: | |
| all_types.append(t) | |
| default_type = all_types[0] if all_types else "Canny" | |
| series_choices = [] | |
| if default_type: | |
| for model in cn_config: | |
| if default_type in model.get("Type", []): | |
| s = model.get("Series", "Default") | |
| if s not in series_choices: | |
| series_choices.append(s) | |
| default_series = series_choices[0] if series_choices else "" | |
| filepath = "" | |
| if default_series and default_type: | |
| for model in cn_config: | |
| if model.get("Series") == default_series and default_type in model.get("Type", []): | |
| filepath = model.get("Filepath", "") | |
| break | |
| return all_types, default_type, series_choices, default_series, filepath | |
| def process_h3_controlnet_inputs(ui_values: dict, prefix: str = "", progress=None) -> list: | |
| active_cns = [] | |
| # 1. Check direct prefix format (e.g. h3_controlnet_videos_h3_fl2va) | |
| videos = ui_values.get(f'h3_controlnet_videos_{prefix}', []) if prefix else [] | |
| types = ui_values.get(f'h3_controlnet_types_{prefix}', []) if prefix else [] | |
| series = ui_values.get(f'h3_controlnet_series_{prefix}', []) if prefix else [] | |
| strengths = ui_values.get(f'h3_controlnet_strengths_{prefix}', []) if prefix else [] | |
| start_percents = ui_values.get(f'h3_controlnet_start_percents_{prefix}', []) if prefix else [] | |
| end_percents = ui_values.get(f'h3_controlnet_end_percents_{prefix}', []) if prefix else [] | |
| filepaths = ui_values.get(f'h3_controlnet_filepaths_{prefix}', []) if prefix else [] | |
| if isinstance(videos, list) and isinstance(types, list): | |
| for idx, vid in enumerate(videos): | |
| if vid: | |
| saved_vid_name = save_temp_video(vid) if isinstance(vid, str) and os.path.exists(vid) else vid | |
| if not saved_vid_name: | |
| continue | |
| fp = filepaths[idx] if idx < len(filepaths) and filepaths[idx] and filepaths[idx] != "None" else None | |
| if not fp: | |
| fp = "minimax_h3_fun_controlnet_union_pruned_int8_convrot.safetensors" | |
| ensure_file_downloaded(fp, progress=progress) | |
| st = float(strengths[idx]) if idx < len(strengths) and strengths[idx] is not None else 1.0 | |
| sp = float(start_percents[idx]) if idx < len(start_percents) and start_percents[idx] is not None else 0.0 | |
| ep = float(end_percents[idx]) if idx < len(end_percents) and end_percents[idx] is not None else 1.0 | |
| active_cns.append({ | |
| "video": saved_vid_name, | |
| "control_net_name": fp, | |
| "strength": st, | |
| "start_percent": sp, | |
| "end_percent": ep | |
| }) | |
| # 2. Check flat component list (e.g. [video1, type1, series1, strength1, start1, end1, filepath1, ...]) | |
| cn_data = ui_values.get(f'h3_controlnet_data_{prefix}', []) or ui_values.get('h3_controlnet_data', []) | |
| if cn_data and not active_cns and isinstance(cn_data, list): | |
| stride = 7 | |
| for i in range(0, len(cn_data), stride): | |
| chunk = cn_data[i:i+stride] | |
| if len(chunk) >= 1 and chunk[0]: | |
| vid = chunk[0] | |
| saved_vid_name = save_temp_video(vid) if isinstance(vid, str) and os.path.exists(vid) else vid | |
| if not saved_vid_name: | |
| continue | |
| fp = chunk[6] if len(chunk) > 6 and chunk[6] and chunk[6] != "None" else "minimax_h3_fun_controlnet_union_pruned_int8_convrot.safetensors" | |
| ensure_file_downloaded(fp, progress=progress) | |
| st = float(chunk[3]) if len(chunk) > 3 and chunk[3] is not None else 1.0 | |
| sp = float(chunk[4]) if len(chunk) > 4 and chunk[4] is not None else 0.0 | |
| ep = float(chunk[5]) if len(chunk) > 5 and chunk[5] is not None else 1.0 | |
| active_cns.append({ | |
| "video": saved_vid_name, | |
| "control_net_name": fp, | |
| "strength": st, | |
| "start_percent": sp, | |
| "end_percent": ep | |
| }) | |
| # 3. Check direct 'h3_controlnets' list of dicts (from MCP or custom payload) | |
| raw_cns = ui_values.get('h3_controlnets', []) | |
| if raw_cns and not active_cns and isinstance(raw_cns, list): | |
| for item in raw_cns: | |
| if isinstance(item, dict): | |
| vid = item.get('video') or item.get('control_video') or item.get('file') | |
| if vid: | |
| saved_vid_name = save_temp_video(vid) if isinstance(vid, str) and os.path.exists(vid) else vid | |
| if not saved_vid_name: | |
| continue | |
| fp = item.get('control_net_name') or item.get('filepath') or item.get('name') or "minimax_h3_fun_controlnet_union_pruned_int8_convrot.safetensors" | |
| ensure_file_downloaded(fp, progress=progress) | |
| active_cns.append({ | |
| "video": saved_vid_name, | |
| "control_net_name": fp, | |
| "strength": float(item.get('strength', 1.0)), | |
| "start_percent": float(item.get('start_percent', 0.0)), | |
| "end_percent": float(item.get('end_percent', 1.0)) | |
| }) | |
| return active_cns | |
| def process_h3_guide_inputs(ui_values: dict, prefix: str = "", progress=None) -> list: | |
| """ | |
| Parses and prepares MiniMax H3 keyframe guide inputs. | |
| Supports: | |
| 1. Direct prefix format (e.g. h3_guide_images_h3_ref2va, h3_guide_times_h3_ref2va, etc.) | |
| 2. Direct 'h3_guides' or 'guides' list of dictionaries (from MCP or custom payload) | |
| Returns: | |
| List of guide dicts: [{'frame_idx': int, 'image': str, 'video': str, 'audio': str}, ...] | |
| """ | |
| active_guides = [] | |
| # 1. Direct prefix format from Gradio UI | |
| images = ui_values.get(f'h3_guide_images_{prefix}', []) if prefix else [] | |
| videos = ui_values.get(f'h3_guide_videos_{prefix}', []) if prefix else [] | |
| audios = ui_values.get(f'h3_guide_audios_{prefix}', []) if prefix else [] | |
| times = ui_values.get(f'h3_guide_times_{prefix}', []) if prefix else [] | |
| frames = ui_values.get(f'h3_guide_frames_{prefix}', []) if prefix else [] | |
| if images or videos or audios: | |
| max_len = max(len(images), len(videos), len(audios), len(times), len(frames)) | |
| for i in range(max_len): | |
| img = images[i] if i < len(images) else None | |
| vid = videos[i] if i < len(videos) else None | |
| aud = audios[i] if i < len(audios) else None | |
| t_val = times[i] if i < len(times) else None | |
| f_val = frames[i] if i < len(frames) else None | |
| saved_img = save_temp_image(img) if img is not None else None | |
| saved_vid = save_temp_video(vid) if vid else None | |
| saved_aud = save_temp_audio(aud) if aud else None | |
| if not (saved_img or saved_vid or saved_aud): | |
| continue | |
| if f_val is not None and str(f_val).strip() != "": | |
| try: | |
| frame_idx = int(round(float(f_val))) | |
| except (ValueError, TypeError): | |
| frame_idx = 0 | |
| elif t_val is not None and str(t_val).strip() != "": | |
| try: | |
| frame_idx = int(round(float(t_val) * 24)) | |
| except (ValueError, TypeError): | |
| frame_idx = 0 | |
| else: | |
| frame_idx = 0 | |
| guide_dict = {"frame_idx": max(0, frame_idx)} | |
| if saved_img: | |
| guide_dict["image"] = saved_img | |
| if saved_vid: | |
| guide_dict["video"] = saved_vid | |
| if saved_aud: | |
| guide_dict["audio"] = saved_aud | |
| active_guides.append(guide_dict) | |
| # 2. Check direct 'h3_guides' or 'guides' list of dicts (from MCP or custom payload) | |
| raw_guides = ui_values.get(f'h3_guides_{prefix}') or ui_values.get('h3_guides') or ui_values.get('guides', []) | |
| if raw_guides and not active_guides and isinstance(raw_guides, list): | |
| for item in raw_guides: | |
| if isinstance(item, dict): | |
| img = item.get('image') | |
| vid = item.get('video') | |
| aud = item.get('audio') | |
| saved_img = save_temp_image(img) if img is not None else None | |
| saved_vid = save_temp_video(vid) if vid else None | |
| saved_aud = save_temp_audio(aud) if aud else None | |
| if not (saved_img or saved_vid or saved_aud): | |
| continue | |
| if item.get('frame_idx') is not None: | |
| try: | |
| frame_idx = int(round(float(item['frame_idx']))) | |
| except (ValueError, TypeError): | |
| frame_idx = 0 | |
| elif item.get('time_seconds') is not None or item.get('time') is not None: | |
| try: | |
| raw_t = item.get('time_seconds') if item.get('time_seconds') is not None else item.get('time') | |
| frame_idx = int(round(float(raw_t) * 24)) | |
| except (ValueError, TypeError): | |
| frame_idx = 0 | |
| else: | |
| frame_idx = 0 | |
| guide_dict = {"frame_idx": max(0, frame_idx)} | |
| if saved_img: | |
| guide_dict["image"] = saved_img | |
| if saved_vid: | |
| guide_dict["video"] = saved_vid | |
| if saved_aud: | |
| guide_dict["audio"] = saved_aud | |
| active_guides.append(guide_dict) | |
| active_guides.sort(key=lambda x: x['frame_idx']) | |
| return active_guides |