maxdemon/CogVideoX-5B-Max-Fun

1import time 2import torch3 4from cogvideox.api.api import infer_forward_api, update_diffusion_transformer_api, update_edition_api5from cogvideox.ui.ui import ui_modelscope, ui_eas, ui6 7if __name__ == "__main__":8    # Choose the ui mode  9    ui_mode = "eas"10    11    # Low gpu memory mode, this is used when the GPU memory is under 16GB12    low_gpu_memory_mode = False13    # Use torch.float16 if GPU does not support torch.bfloat1614    # ome graphics cards, such as v100, 2080ti, do not support torch.bfloat1615    weight_dtype = torch.bfloat1616 17    # Server ip18    server_name = "0.0.0.0"19    server_port = 786020 21    # Params below is used when ui_mode = "modelscope"22    model_name = "models/Diffusion_Transformer/CogVideoX-Fun-V1.1-5b-InP"23    # "Inpaint" or "Control"24    model_type = "Inpaint"25    # Save dir of this model26    savedir_sample = "samples"27 28    if ui_mode == "modelscope":29        demo, controller = ui_modelscope(model_name, model_type, savedir_sample, low_gpu_memory_mode, weight_dtype)30    elif ui_mode == "eas":31        demo, controller = ui_eas(model_name, savedir_sample)32    else:33        demo, controller = ui(low_gpu_memory_mode, weight_dtype)34 35    # launch gradio36    app, _, _ = demo.queue(status_update_rate=1).launch(37        server_name=server_name,38        server_port=server_port,39        prevent_thread_lock=True40    )41    42    # launch api43    infer_forward_api(None, app, controller)44    update_diffusion_transformer_api(None, app, controller)45    update_edition_api(None, app, controller)46    47    # not close the python48    while True:49        time.sleep(5)