Class MultiNodesEngine

videox_fun/api/api_multi_nodes.py:156–197

                raise HTTPException(status_code=500, detail=str(e))

    class MultiNodesEngine:
        def __init__(
            self, 
            world_size, 
            Controller,
            GPU_memory_mode, 
            scheduler_dict, 
            model_name, 
            model_type, 
            config_path,
            ulysses_degree, 
            ring_degree, 
            enable_teacache, 
            teacache_threshold, 
            num_skip_start_steps, 
            teacache_offload, 
            weight_dtype,
            savedir_sample
        ):
            # Ensure Ray is initialized
            if not ray.is_initialized():
                ray.init()
            
            num_workers = world_size
            self.workers = [
                MultiNodesGenerator.remote(
                    rank, world_size, Controller, 
                    GPU_memory_mode, scheduler_dict, model_name=model_name, model_type=model_type, config_path=config_path, 
                    ulysses_degree=ulysses_degree, ring_degree=ring_degree, enable_teacache=enable_teacache, teacache_threshold=teacache_threshold, num_skip_start_steps=num_skip_start_steps, 
                    teacache_offload=teacache_offload, weight_dtype=weight_dtype, savedir_sample=savedir_sample,
                )
                for rank in range(num_workers)
            ]
            print("Update workers done")
            
        async def generate(self, data):
            results = ray.get([
                worker.generate.remote(data)
                for worker in self.workers
            ])

            return next(path for path in results if path is not None) 

    def multi_nodes_infer_forward_api(_: gr.Blocks, app: FastAPI, engine):