Spaces:
Running on Zero
Running on Zero
| import os | |
| import glob | |
| from typing import Any, Dict, List | |
| import gradio as gr | |
| import spaces | |
| from core.pipelines.workflow_executor import WorkflowExecutor | |
| from core.model_manager import model_manager | |
| def _extract_required_models(workflow: Dict[str, Any]) -> List[str]: | |
| required_models = [] | |
| MODEL_KEYS = {'unet_name', 'clip_name', 'vae_name', 'model_name', 'ckpt_name', 'clip_name1', 'clip_name2'} | |
| MODEL_EXTENSIONS = ('.safetensors', '.ckpt', '.pt', '.bin', '.pth') | |
| for node_id, node_data in workflow.items(): | |
| if node_data.get('class_type') == 'LoraLoader': | |
| continue | |
| inputs = node_data.get('inputs', {}) | |
| for k, v in inputs.items(): | |
| if isinstance(v, str) and v and k != 'lora_name': | |
| if k in MODEL_KEYS or any(v.lower().endswith(ext) for ext in MODEL_EXTENSIONS): | |
| if v not in required_models and not v.startswith('[') and not v.endswith(']'): | |
| required_models.append(v) | |
| return required_models | |
| def _extract_video_paths(out: Any) -> List[str]: | |
| video_files = [] | |
| def process_item(item): | |
| if item is None: | |
| return | |
| if isinstance(item, (list, tuple)): | |
| for sub in item: | |
| process_item(sub) | |
| return | |
| if isinstance(item, str) and os.path.exists(item): | |
| video_files.append(item) | |
| return | |
| for attr in ['path', 'saved_path', 'filepath', 'filename', 'full_path']: | |
| val = getattr(item, attr, None) | |
| if isinstance(val, str) and os.path.exists(val): | |
| video_files.append(val) | |
| return | |
| elif isinstance(val, str): | |
| try: | |
| import folder_paths | |
| possible = os.path.join(folder_paths.get_output_directory(), val) | |
| if os.path.exists(possible): | |
| video_files.append(possible) | |
| return | |
| except Exception: | |
| pass | |
| if hasattr(item, '__dict__'): | |
| for k, v in item.__dict__.items(): | |
| if isinstance(v, str) and (v.endswith('.mp4') or v.endswith('.webm') or v.endswith('.gif')): | |
| if os.path.exists(v): | |
| video_files.append(v) | |
| return | |
| try: | |
| import folder_paths | |
| possible = os.path.join(folder_paths.get_output_directory(), v) | |
| if os.path.exists(possible): | |
| video_files.append(possible) | |
| return | |
| except Exception: | |
| pass | |
| if isinstance(item, dict): | |
| ui_info = item.get("ui", {}) | |
| for key in ["images", "videos", "video"]: | |
| for sub in ui_info.get(key, []): | |
| if isinstance(sub, dict) and "filename" in sub: | |
| fn = sub["filename"] | |
| subfolder = sub.get("subfolder", "") | |
| try: | |
| import folder_paths | |
| full = os.path.join(folder_paths.get_output_directory(), subfolder, fn) if subfolder else os.path.join(folder_paths.get_output_directory(), fn) | |
| if os.path.exists(full): | |
| video_files.append(full) | |
| return | |
| except Exception: | |
| pass | |
| process_item(out) | |
| if not video_files: | |
| try: | |
| import folder_paths | |
| out_dir = folder_paths.get_output_directory() | |
| mp4_files = glob.glob(os.path.join(out_dir, "**", "*.mp4"), recursive=True) | |
| if mp4_files: | |
| mp4_files.sort(key=os.path.getmtime, reverse=True) | |
| video_files.append(mp4_files[0]) | |
| except Exception: | |
| pass | |
| return video_files | |
| def _execute_h3_workflow_gpu(workflow: Dict[str, Any]): | |
| initial_objects = {} | |
| return WorkflowExecutor.execute_workflow(workflow, initial_objects=initial_objects) | |
| def generate_h3_video_wrapper(process_inputs_func, ui_inputs: dict, progress=gr.Progress(track_tqdm=True)): | |
| progress(0.1, desc="Assembling H3 Video Workflow...") | |
| try: | |
| workflow, extra_data = process_inputs_func(ui_inputs, progress=progress) | |
| except TypeError: | |
| workflow, extra_data = process_inputs_func(ui_inputs) | |
| required_models = _extract_required_models(workflow) | |
| if required_models: | |
| progress(0.2, desc=f"Ensuring models are downloaded: {len(required_models)} file(s)...") | |
| model_manager.ensure_models_downloaded(required_models, progress=progress) | |
| zero_gpu_duration = ui_inputs.get('zero_gpu_duration', 120) | |
| try: | |
| duration = int(zero_gpu_duration) | |
| except (ValueError, TypeError): | |
| duration = 120 | |
| if duration <= 0: | |
| duration = 120 | |
| progress(0.4, desc=f"Executing H3 Video Workflow on GPU (ZeroGPU {duration}s)...") | |
| gpu_runner = spaces.GPU(duration=duration)(_execute_h3_workflow_gpu) | |
| raw_output = gpu_runner(workflow) | |
| video_files = _extract_video_paths(raw_output) | |
| final_video = video_files[0] if video_files else None | |
| print("✅ H3 Video generated successfully") | |
| progress(1.0, desc="H3 Video Generation Complete!") | |
| return final_video |