Spaces:
Runtime error
Runtime error
| import os | |
| import sys | |
| import unittest | |
| from unittest.mock import MagicMock | |
| # Add project root to sys.path | |
| PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) | |
| if PROJECT_ROOT not in sys.path: | |
| sys.path.insert(0, PROJECT_ROOT) | |
| import asyncio | |
| if 'torch' not in sys.modules: | |
| sys.modules['torch'] = MagicMock() | |
| if 'numpy' not in sys.modules: | |
| sys.modules['numpy'] = MagicMock() | |
| if 'gradio' not in sys.modules: | |
| sys.modules['gradio'] = MagicMock() | |
| if 'spaces' not in sys.modules: | |
| mock_spaces = MagicMock() | |
| mock_spaces.GPU = lambda fn: fn | |
| sys.modules['spaces'] = mock_spaces | |
| if 'huggingface_hub' not in sys.modules: | |
| mock_hf = MagicMock() | |
| sys.modules['huggingface_hub'] = mock_hf | |
| sys.modules['huggingface_hub.constants'] = MagicMock() | |
| if 'PIL' not in sys.modules: | |
| mock_pil = MagicMock() | |
| class MockImage: | |
| pass | |
| mock_pil.Image = MagicMock() | |
| mock_pil.Image.Image = MockImage | |
| sys.modules['PIL'] = mock_pil | |
| sys.modules['PIL.Image'] = mock_pil.Image | |
| # Mock ComfyUI environment if not installed locally | |
| if 'execution' not in sys.modules: | |
| sys.modules['execution'] = MagicMock() | |
| if 'server' not in sys.modules: | |
| sys.modules['server'] = MagicMock() | |
| if 'nodes' not in sys.modules: | |
| mock_nodes = MagicMock() | |
| async def async_noop(): | |
| pass | |
| mock_nodes.init_extra_nodes = async_noop | |
| dummy_cls = MagicMock() | |
| mock_nodes.NODE_CLASS_MAPPINGS = { | |
| 'CLIPTextEncode': dummy_cls, | |
| 'CLIPTextEncodeSDXL': dummy_cls, | |
| 'LoraLoader': dummy_cls, | |
| 'CLIPSetLastLayer': dummy_cls, | |
| 'CheckpointLoaderSimple': dummy_cls, | |
| } | |
| mock_nodes.CheckpointLoaderSimple = dummy_cls | |
| mock_nodes.LoraLoader = dummy_cls | |
| sys.modules['nodes'] = mock_nodes | |
| # Helper to create dummy ComfyUI node class | |
| def make_mock_node(name, required_inputs=None, optional_inputs=None): | |
| req = required_inputs or {} | |
| opt = optional_inputs or {} | |
| class MockNode: | |
| NODE_NAME = name | |
| def INPUT_TYPES(cls): | |
| return {"required": req, "optional": opt} | |
| return MockNode | |
| from comfy_integration import nodes as ci_nodes | |
| # Register mock nodes for testing | |
| TEST_NODE_MAPPINGS = { | |
| "UNETLoader": make_mock_node("UNETLoader", {"unet_name": (["model.safetensors"],), "weight_dtype": (["default"],)}), | |
| "CLIPLoader": make_mock_node("CLIPLoader", {"clip_name": (["clip.safetensors"],), "type": (["minimax"],)}), | |
| "VAELoader": make_mock_node("VAELoader", {"vae_name": (["vae.safetensors"],)}), | |
| "MiniMaxH3ReferenceToVideo": make_mock_node("MiniMaxH3ReferenceToVideo", { | |
| "clip": ("CLIP",), "vae": ("VAE",), "audio_vae": ("VAE",), | |
| "prompt": ("STRING", {"default": ""}), "width": ("INT", {"default": 960}), | |
| "height": ("INT", {"default": 544}), "length": ("INT", {"default": 73}) | |
| }), | |
| "BasicGuider": make_mock_node("BasicGuider", {"model": ("MODEL",), "conditioning": ("CONDITIONING",)}), | |
| "KSamplerSelect": make_mock_node("KSamplerSelect", {"sampler_name": (["res_multistep"], {"default": "res_multistep"})}), | |
| "BasicScheduler": make_mock_node("BasicScheduler", {"model": ("MODEL",), "scheduler": (["simple"], {"default": "simple"}), "steps": ("INT", {"default": 20})}), | |
| "RandomNoise": make_mock_node("RandomNoise", {"noise_seed": ("INT", {"default": 0})}), | |
| "SamplerCustomAdvanced": make_mock_node("SamplerCustomAdvanced", { | |
| "noise": ("NOISE",), "guider": ("GUIDER",), "sampler": ("SAMPLER",), | |
| "sigmas": ("SIGMAS",), "latent_image": ("LATENT",) | |
| }), | |
| "VAEDecode": make_mock_node("VAEDecode", {"samples": ("LATENT",), "vae": ("VAE",)}), | |
| "VAEDecodeAudio": make_mock_node("VAEDecodeAudio", {"samples": ("LATENT",), "vae": ("VAE",)}), | |
| "CreateVideo": make_mock_node("CreateVideo", {"images": ("IMAGE",), "audio": ("AUDIO",), "fps": ("INT", {"default": 24})}), | |
| "SaveVideo": make_mock_node("SaveVideo", {"video": ("VIDEO",), "filename_prefix": ("STRING", {"default": "video/test"})}), | |
| "LoadImage": make_mock_node("LoadImage", {"image": ("STRING", {"default": ""})}), | |
| "LoadVideo": make_mock_node("LoadVideo", {"file": ("STRING", {"default": ""})}), | |
| "GetVideoComponents": make_mock_node("GetVideoComponents", {"video": ("VIDEO",)}), | |
| "LoadAudio": make_mock_node("LoadAudio", {"audio": ("STRING", {"default": ""})}), | |
| "MiniMaxH3AddGuide": make_mock_node("MiniMaxH3AddGuide", { | |
| "positive": ("CONDITIONING",), | |
| "latent": ("LATENT",), | |
| "frame_idx": ("INT", {"default": 0}), | |
| }, { | |
| "vae": ("VAE",), | |
| "audio_vae": ("VAE",), | |
| "image": ("IMAGE",), | |
| "audio": ("AUDIO",), | |
| }), | |
| "LoraLoader": make_mock_node("LoraLoader", {"model": ("MODEL",), "clip": ("CLIP",), "lora_name": ("STRING",)}), | |
| "MiniMaxH3ImageToVideo": make_mock_node("MiniMaxH3ImageToVideo", { | |
| "clip": ("CLIP",), "vae": ("VAE",), | |
| "prompt": ("STRING", {"default": ""}), "width": ("INT", {"default": 1344}), | |
| "height": ("INT", {"default": 768}), "length": ("INT", {"default": 73}) | |
| }, { | |
| "first_frame": ("IMAGE",), "last_frame": ("IMAGE",) | |
| }), | |
| "ImageScale": make_mock_node("ImageScale", { | |
| "image": ("IMAGE",), "width": ("INT", {"default": 512}), "height": ("INT", {"default": 512}), | |
| "upscale_method": (["nearest-exact"],), "crop": (["center"],) | |
| }), | |
| } | |
| ci_nodes.NODE_CLASS_MAPPINGS.update(TEST_NODE_MAPPINGS) | |
| from chain_injectors import discover_injectors, get_registered_features | |
| from utils.app_utils import process_h3_guide_inputs | |
| from core.workflow_assembler import WorkflowAssembler | |
| class TestH3GuideIntegration(unittest.TestCase): | |
| def setUp(self): | |
| self.recipe_path = os.path.join(PROJECT_ROOT, "ui", "shared", "h3_ref2va_recipe.yaml") | |
| def test_injector_discovery(self): | |
| injectors = discover_injectors() | |
| self.assertIn("dynamic_h3_guide_chains", injectors) | |
| self.assertTrue(callable(injectors["dynamic_h3_guide_chains"])) | |
| def test_registered_features(self): | |
| import yaml | |
| features = get_registered_features() | |
| self.assertIn("h3_guide", features) | |
| self.assertEqual(features["h3_guide"]["chain_type"], "dynamic_h3_guide_chains") | |
| with open(os.path.join(PROJECT_ROOT, "yaml", "chain_features.yaml"), "r", encoding="utf-8") as f: | |
| chain_features = yaml.safe_load(f) | |
| self.assertIn("h3_guide", chain_features) | |
| schema = chain_features["h3_guide"]["parameters_schema"] | |
| self.assertIn("frame_idx", schema["properties"]) | |
| self.assertIn("time_seconds", schema["properties"]) | |
| self.assertIn("image", schema["properties"]) | |
| self.assertIn("video", schema["properties"]) | |
| self.assertIn("audio", schema["properties"]) | |
| def test_process_h3_guide_inputs_direct_list(self): | |
| raw_guides = [ | |
| {"image": "path/to/img1.png", "time_seconds": 1.5}, | |
| {"image": "path/to/img2.png", "frame_idx": 10}, | |
| ] | |
| ui_values = {"h3_guides": raw_guides} | |
| processed = process_h3_guide_inputs(ui_values) | |
| self.assertEqual(len(processed), 2) | |
| # Should be sorted by frame_idx: 10 comes before round(1.5 * 24) = 36 | |
| self.assertEqual(processed[0]["frame_idx"], 10) | |
| self.assertEqual(processed[1]["frame_idx"], 36) | |
| self.assertEqual(processed[0]["image"], "img2.png") | |
| self.assertEqual(processed[1]["image"], "img1.png") | |
| def test_process_h3_guide_inputs_ui_prefix(self): | |
| ui_values = { | |
| "h3_guide_images_test": ["img_a.png", "img_b.png"], | |
| "h3_guide_videos_test": [None, None], | |
| "h3_guide_audios_test": [None, None], | |
| "h3_guide_times_test": [2.0, 0.5], | |
| "h3_guide_frames_test": [None, None], | |
| } | |
| processed = process_h3_guide_inputs(ui_values, prefix="test") | |
| self.assertEqual(len(processed), 2) | |
| # Sorted: 0.5s (12 frames) then 2.0s (48 frames) | |
| self.assertEqual(processed[0]["frame_idx"], 12) | |
| self.assertEqual(processed[1]["frame_idx"], 48) | |
| self.assertEqual(processed[0]["image"], "img_b.png") | |
| self.assertEqual(processed[1]["image"], "img_a.png") | |
| def test_workflow_assembler_without_guides(self): | |
| assembler = WorkflowAssembler(self.recipe_path) | |
| ui_values = { | |
| "prompt": "test prompt", | |
| "width": 960, | |
| "height": 544, | |
| "length": 73, | |
| "h3_guides": [] | |
| } | |
| workflow = assembler.assemble(ui_values) | |
| # BasicGuider conditioning should connect directly to minimax_h3:0 | |
| guider_id = assembler.node_map["basic_guider"] | |
| minimax_id = assembler.node_map["minimax_h3"] | |
| cond_input = workflow[guider_id]["inputs"]["conditioning"] | |
| self.assertEqual(cond_input, [minimax_id, 0]) | |
| def test_workflow_assembler_with_single_guide(self): | |
| assembler = WorkflowAssembler(self.recipe_path) | |
| ui_values = { | |
| "prompt": "test prompt", | |
| "width": 960, | |
| "height": 544, | |
| "length": 73, | |
| "h3_guides": [ | |
| {"image": "test_guide.png", "frame_idx": 24} | |
| ] | |
| } | |
| workflow = assembler.assemble(ui_values) | |
| guider_id = assembler.node_map["basic_guider"] | |
| minimax_id = assembler.node_map["minimax_h3"] | |
| vae_id = assembler.node_map["video_vae_loader"] | |
| audio_vae_id = assembler.node_map["audio_vae_loader"] | |
| # Find MiniMaxH3AddGuide node in workflow | |
| guide_nodes = [ | |
| (nid, n) for nid, n in workflow.items() | |
| if n.get("class_type") == "MiniMaxH3AddGuide" | |
| ] | |
| self.assertEqual(len(guide_nodes), 1) | |
| guide_id, guide_data = guide_nodes[0] | |
| # Verify guide inputs | |
| self.assertEqual(guide_data["inputs"]["frame_idx"], 24) | |
| self.assertEqual(guide_data["inputs"]["positive"], [minimax_id, 0]) | |
| self.assertEqual(guide_data["inputs"]["latent"], [minimax_id, 1]) | |
| self.assertEqual(guide_data["inputs"]["vae"], [vae_id, 0]) | |
| self.assertEqual(guide_data["inputs"]["audio_vae"], [audio_vae_id, 0]) | |
| # Verify basic_guider points to the guide output | |
| self.assertEqual(workflow[guider_id]["inputs"]["conditioning"], [guide_id, 0]) | |
| def test_workflow_assembler_with_chained_guides(self): | |
| assembler = WorkflowAssembler(self.recipe_path) | |
| ui_values = { | |
| "prompt": "test prompt", | |
| "width": 960, | |
| "height": 544, | |
| "length": 73, | |
| "h3_guides": [ | |
| {"image": "guide1.png", "frame_idx": 10}, | |
| {"image": "guide2.png", "frame_idx": 36}, | |
| {"image": "guide3.png", "frame_idx": 72} | |
| ] | |
| } | |
| workflow = assembler.assemble(ui_values) | |
| guider_id = assembler.node_map["basic_guider"] | |
| minimax_id = assembler.node_map["minimax_h3"] | |
| guide_nodes = [ | |
| (nid, n) for nid, n in workflow.items() | |
| if n.get("class_type") == "MiniMaxH3AddGuide" | |
| ] | |
| self.assertEqual(len(guide_nodes), 3) | |
| # Sort by frame_idx | |
| guide_nodes_sorted = sorted(guide_nodes, key=lambda x: x[1]["inputs"]["frame_idx"]) | |
| g1_id, g1 = guide_nodes_sorted[0] | |
| g2_id, g2 = guide_nodes_sorted[1] | |
| g3_id, g3 = guide_nodes_sorted[2] | |
| self.assertEqual(g1["inputs"]["positive"], [minimax_id, 0]) | |
| self.assertEqual(g2["inputs"]["positive"], [g1_id, 0]) | |
| self.assertEqual(g3["inputs"]["positive"], [g2_id, 0]) | |
| # BasicGuider connects to last guide in chain | |
| self.assertEqual(workflow[guider_id]["inputs"]["conditioning"], [g3_id, 0]) | |
| def test_workflow_assembler_with_video_and_audio_guides(self): | |
| assembler = WorkflowAssembler(self.recipe_path) | |
| ui_values = { | |
| "prompt": "test prompt", | |
| "width": 960, | |
| "height": 544, | |
| "length": 73, | |
| "h3_guides": [ | |
| {"video": "clip.mp4", "frame_idx": 24}, | |
| {"audio": "audio.mp3", "frame_idx": 48} | |
| ] | |
| } | |
| workflow = assembler.assemble(ui_values) | |
| guide_nodes = [ | |
| (nid, n) for nid, n in workflow.items() | |
| if n.get("class_type") == "MiniMaxH3AddGuide" | |
| ] | |
| self.assertEqual(len(guide_nodes), 2) | |
| # Check LoadVideo + GetVideoComponents created | |
| load_vids = [n for n in workflow.values() if n.get("class_type") == "LoadVideo"] | |
| get_comps = [n for n in workflow.values() if n.get("class_type") == "GetVideoComponents"] | |
| load_auds = [n for n in workflow.values() if n.get("class_type") == "LoadAudio"] | |
| self.assertEqual(len(load_vids), 1) | |
| self.assertEqual(len(get_comps), 1) | |
| self.assertEqual(len(load_auds), 1) | |
| def test_h3_guide_ui_configuration(self): | |
| from ui.shared.ui_components import create_h3_guide_ui | |
| import gradio as gr | |
| gr.Slider.reset_mock() | |
| gr.Number.reset_mock() | |
| comps = create_h3_guide_ui("test_ui") | |
| times = comps.get("h3_guide_times_test_ui", []) | |
| frames = comps.get("h3_guide_frames_test_ui", []) | |
| # Step must be 0.5 for all slider calls | |
| self.assertGreater(len(times), 0) | |
| self.assertGreater(len(gr.Slider.call_args_list), 0) | |
| for call in gr.Slider.call_args_list: | |
| self.assertEqual(call.kwargs.get("step"), 0.5) | |
| # No Frame Index components displayed | |
| self.assertEqual(len(frames), 0) | |
| for call in gr.Number.call_args_list: | |
| label = call.kwargs.get("label", "") | |
| self.assertNotIn("Frame Index", label) | |
| self.assertNotIn("24 fps", label.lower()) | |
| def test_h3_fl2va_recipe_guide_integration(self): | |
| fl2va_recipe_path = os.path.join(PROJECT_ROOT, "ui", "shared", "h3_fl2va_recipe.yaml") | |
| assembler = WorkflowAssembler(fl2va_recipe_path) | |
| ui_values = { | |
| "prompt": "sunset test", | |
| "width": 1344, | |
| "height": 768, | |
| "length": 73, | |
| "h3_guides": [ | |
| {"image": "intermediate_keyframe.png", "frame_idx": 36} | |
| ] | |
| } | |
| workflow = assembler.assemble(ui_values) | |
| guider_id = assembler.node_map["basic_guider"] | |
| minimax_id = assembler.node_map["minimax_h3"] | |
| guide_nodes = [ | |
| (nid, n) for nid, n in workflow.items() | |
| if n.get("class_type") == "MiniMaxH3AddGuide" | |
| ] | |
| self.assertEqual(len(guide_nodes), 1) | |
| guide_id, guide_data = guide_nodes[0] | |
| self.assertEqual(guide_data["inputs"]["frame_idx"], 36) | |
| self.assertEqual(guide_data["inputs"]["positive"], [minimax_id, 0]) | |
| self.assertEqual(workflow[guider_id]["inputs"]["conditioning"], [guide_id, 0]) | |
| def test_h3_fl2va_ui_integration(self): | |
| from ui.shared import h3_fl2va_ui | |
| comps = h3_fl2va_ui.create_ui() | |
| self.assertIn("h3_guide_times_h3_fl2va", comps) | |
| self.assertIn("h3_guide_images_h3_fl2va", comps) | |
| self.assertIn("h3_guide_videos_h3_fl2va", comps) | |
| self.assertIn("h3_guide_audios_h3_fl2va", comps) | |
| def test_unified_h3_ui_integration(self): | |
| from ui.h3_ui import create_ui | |
| comps = create_ui() | |
| self.assertIn("task", comps) | |
| self.assertIn("first_frame", comps) | |
| self.assertIn("last_frame", comps) | |
| self.assertIn("h3_guide_times_h3", comps) | |
| self.assertIn("h3_guide_images_h3", comps) | |
| self.assertIn("h3_guide_videos_h3", comps) | |
| self.assertIn("h3_guide_audios_h3", comps) | |
| if __name__ == "__main__": | |
| unittest.main() | |