Bubbliiiing
1a494b9188
Update Wan2.2 Comfyui, fix bug in api, fix bug in dataloader, fix bug in readme, fix bug in multigpu
320 lines
14 KiB
Python
Executable File
320 lines
14 KiB
Python
Executable File
# This file is modified from https://github.com/xdit-project/xDiT/blob/main/entrypoints/launch.py
|
|
import base64
|
|
import gc
|
|
import hashlib
|
|
import io
|
|
import os
|
|
import tempfile
|
|
from io import BytesIO
|
|
|
|
import gradio as gr
|
|
import requests
|
|
import torch
|
|
import torch.distributed as dist
|
|
from fastapi import FastAPI, HTTPException
|
|
from PIL import Image
|
|
|
|
from .api import download_from_url, encode_file_to_base64
|
|
|
|
try:
|
|
import ray
|
|
except:
|
|
print("Ray is not installed. If you want to use multi gpus api. Please install it by running 'pip install ray'.")
|
|
ray = None
|
|
|
|
def save_base64_video_dist(base64_string):
|
|
video_data = base64.b64decode(base64_string)
|
|
|
|
md5_hash = hashlib.md5(video_data).hexdigest()
|
|
filename = f"{md5_hash}.mp4"
|
|
|
|
temp_dir = tempfile.gettempdir()
|
|
file_path = os.path.join(temp_dir, filename)
|
|
|
|
if dist.is_initialized():
|
|
if dist.get_rank() == 0:
|
|
with open(file_path, 'wb') as video_file:
|
|
video_file.write(video_data)
|
|
dist.barrier()
|
|
else:
|
|
with open(file_path, 'wb') as video_file:
|
|
video_file.write(video_data)
|
|
return file_path
|
|
|
|
def save_base64_image_dist(base64_string):
|
|
video_data = base64.b64decode(base64_string)
|
|
|
|
md5_hash = hashlib.md5(video_data).hexdigest()
|
|
filename = f"{md5_hash}.jpg"
|
|
|
|
temp_dir = tempfile.gettempdir()
|
|
file_path = os.path.join(temp_dir, filename)
|
|
|
|
if dist.is_initialized():
|
|
if dist.get_rank() == 0:
|
|
with open(file_path, 'wb') as video_file:
|
|
video_file.write(video_data)
|
|
dist.barrier()
|
|
else:
|
|
with open(file_path, 'wb') as video_file:
|
|
video_file.write(video_data)
|
|
return file_path
|
|
|
|
def save_url_video_dist(url):
|
|
video_data = download_from_url(url)
|
|
if video_data:
|
|
return save_base64_video_dist(base64.b64encode(video_data))
|
|
return None
|
|
|
|
def save_url_image_dist(url):
|
|
image_data = download_from_url(url)
|
|
if image_data:
|
|
return save_base64_image_dist(base64.b64encode(image_data))
|
|
return None
|
|
|
|
if ray is not None:
|
|
@ray.remote(num_gpus=1)
|
|
class MultiNodesGenerator:
|
|
def __init__(
|
|
self, rank: int, world_size: int, Controller,
|
|
GPU_memory_mode, scheduler_dict, model_name=None, model_type="Inpaint",
|
|
config_path=None, ulysses_degree=1, ring_degree=1,
|
|
fsdp_dit=False, fsdp_text_encoder=False, compile_dit=False,
|
|
weight_dtype=None, savedir_sample=None,
|
|
):
|
|
# Set PyTorch distributed environment variables
|
|
os.environ["RANK"] = str(rank)
|
|
os.environ["WORLD_SIZE"] = str(world_size)
|
|
os.environ["MASTER_ADDR"] = "127.0.0.1"
|
|
os.environ["MASTER_PORT"] = "29500"
|
|
|
|
self.rank = rank
|
|
self.controller = 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,
|
|
fsdp_dit=fsdp_dit, fsdp_text_encoder=fsdp_text_encoder, compile_dit=compile_dit,
|
|
weight_dtype=weight_dtype, savedir_sample=savedir_sample,
|
|
)
|
|
|
|
def generate(self, datas):
|
|
try:
|
|
base_model_path = datas.get('base_model_path', 'none')
|
|
base_model_2_path = datas.get('base_model_2_path', 'none')
|
|
lora_model_path = datas.get('lora_model_path', 'none')
|
|
lora_model_2_path = datas.get('lora_model_2_path', 'none')
|
|
lora_alpha_slider = datas.get('lora_alpha_slider', 0.55)
|
|
prompt_textbox = datas.get('prompt_textbox', None)
|
|
negative_prompt_textbox = datas.get('negative_prompt_textbox', 'The video is not of a high quality, it has a low resolution. Watermark present in each frame. The background is solid. Strange body and strange trajectory. Distortion. ')
|
|
sampler_dropdown = datas.get('sampler_dropdown', 'Euler')
|
|
sample_step_slider = datas.get('sample_step_slider', 30)
|
|
resize_method = datas.get('resize_method', "Generate by")
|
|
width_slider = datas.get('width_slider', 672)
|
|
height_slider = datas.get('height_slider', 384)
|
|
base_resolution = datas.get('base_resolution', 512)
|
|
is_image = datas.get('is_image', False)
|
|
generation_method = datas.get('generation_method', False)
|
|
length_slider = datas.get('length_slider', 49)
|
|
overlap_video_length = datas.get('overlap_video_length', 4)
|
|
partial_video_length = datas.get('partial_video_length', 72)
|
|
cfg_scale_slider = datas.get('cfg_scale_slider', 6)
|
|
start_image = datas.get('start_image', None)
|
|
end_image = datas.get('end_image', None)
|
|
validation_video = datas.get('validation_video', None)
|
|
validation_video_mask = datas.get('validation_video_mask', None)
|
|
control_video = datas.get('control_video', None)
|
|
denoise_strength = datas.get('denoise_strength', 0.70)
|
|
seed_textbox = datas.get("seed_textbox", 43)
|
|
|
|
ref_image = datas.get('ref_image', None)
|
|
enable_teacache = datas.get('enable_teacache', True)
|
|
teacache_threshold = datas.get('teacache_threshold', 0.10)
|
|
num_skip_start_steps = datas.get('num_skip_start_steps', 1)
|
|
teacache_offload = datas.get('teacache_offload', False)
|
|
cfg_skip_ratio = datas.get('cfg_skip_ratio', 0)
|
|
enable_riflex = datas.get('enable_riflex', False)
|
|
riflex_k = datas.get('riflex_k', 6)
|
|
fps = datas.get('fps', None)
|
|
|
|
generation_method = "Image Generation" if is_image else generation_method
|
|
|
|
if start_image is not None:
|
|
if start_image.startswith('http'):
|
|
start_image = save_url_image_dist(start_image)
|
|
start_image = [Image.open(start_image).convert("RGB")]
|
|
else:
|
|
start_image = base64.b64decode(start_image)
|
|
start_image = [Image.open(BytesIO(start_image)).convert("RGB")]
|
|
|
|
if end_image is not None:
|
|
if end_image.startswith('http'):
|
|
end_image = save_url_image_dist(end_image)
|
|
end_image = [Image.open(end_image).convert("RGB")]
|
|
else:
|
|
end_image = base64.b64decode(end_image)
|
|
end_image = [Image.open(BytesIO(end_image)).convert("RGB")]
|
|
|
|
if validation_video is not None:
|
|
if validation_video.startswith('http'):
|
|
validation_video = save_url_video_dist(validation_video)
|
|
else:
|
|
validation_video = save_base64_video_dist(validation_video)
|
|
|
|
if validation_video_mask is not None:
|
|
if validation_video_mask.startswith('http'):
|
|
validation_video_mask = save_url_image_dist(validation_video_mask)
|
|
else:
|
|
validation_video_mask = save_base64_image_dist(validation_video_mask)
|
|
|
|
if control_video is not None:
|
|
if control_video.startswith('http'):
|
|
control_video = save_url_video_dist(control_video)
|
|
else:
|
|
control_video = save_base64_video_dist(control_video)
|
|
|
|
if ref_image is not None:
|
|
if ref_image.startswith('http'):
|
|
ref_image = save_url_image_dist(ref_image)
|
|
ref_image = [Image.open(ref_image).convert("RGB")]
|
|
else:
|
|
ref_image = base64.b64decode(ref_image)
|
|
ref_image = [Image.open(BytesIO(ref_image)).convert("RGB")]
|
|
|
|
try:
|
|
save_sample_path, comment = self.controller.generate(
|
|
"",
|
|
base_model_path,
|
|
lora_model_path,
|
|
lora_alpha_slider,
|
|
prompt_textbox,
|
|
negative_prompt_textbox,
|
|
sampler_dropdown,
|
|
sample_step_slider,
|
|
resize_method,
|
|
width_slider,
|
|
height_slider,
|
|
base_resolution,
|
|
generation_method,
|
|
length_slider,
|
|
overlap_video_length,
|
|
partial_video_length,
|
|
cfg_scale_slider,
|
|
start_image,
|
|
end_image,
|
|
validation_video,
|
|
validation_video_mask,
|
|
control_video,
|
|
denoise_strength,
|
|
seed_textbox,
|
|
ref_image = ref_image,
|
|
enable_teacache = enable_teacache,
|
|
teacache_threshold = teacache_threshold,
|
|
num_skip_start_steps = num_skip_start_steps,
|
|
teacache_offload = teacache_offload,
|
|
cfg_skip_ratio = cfg_skip_ratio,
|
|
enable_riflex = enable_riflex,
|
|
riflex_k = riflex_k,
|
|
base_model_2_dropdown = base_model_2_path,
|
|
lora_model_2_dropdown = lora_model_2_path,
|
|
fps = fps,
|
|
is_api = True,
|
|
)
|
|
except Exception as e:
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.ipc_collect()
|
|
save_sample_path = ""
|
|
comment = f"Error. error information is {str(e)}"
|
|
if dist.is_initialized():
|
|
if dist.get_rank() == 0:
|
|
return {"message": comment, "save_sample_path": None, "base64_encoding": None}
|
|
else:
|
|
return None
|
|
else:
|
|
return {"message": comment, "save_sample_path": None, "base64_encoding": None}
|
|
|
|
|
|
if dist.is_initialized():
|
|
if dist.get_rank() == 0:
|
|
if save_sample_path != "":
|
|
return {"message": comment, "save_sample_path": save_sample_path, "base64_encoding": encode_file_to_base64(save_sample_path)}
|
|
else:
|
|
return {"message": comment, "save_sample_path": None, "base64_encoding": None}
|
|
else:
|
|
return None
|
|
else:
|
|
if save_sample_path != "":
|
|
return {"message": comment, "save_sample_path": save_sample_path, "base64_encoding": encode_file_to_base64(save_sample_path)}
|
|
else:
|
|
return {"message": comment, "save_sample_path": None, "base64_encoding": None}
|
|
|
|
except Exception as e:
|
|
print(f"Error generating: {str(e)}")
|
|
comment = f"Error generating: {str(e)}"
|
|
if dist.is_initialized():
|
|
if dist.get_rank() == 0:
|
|
return {"message": comment, "save_sample_path": None, "base64_encoding": None}
|
|
else:
|
|
return None
|
|
else:
|
|
return {"message": comment, "save_sample_path": None, "base64_encoding": None}
|
|
|
|
class MultiNodesEngine:
|
|
def __init__(
|
|
self,
|
|
world_size,
|
|
Controller,
|
|
GPU_memory_mode,
|
|
scheduler_dict,
|
|
model_name,
|
|
model_type,
|
|
config_path,
|
|
ulysses_degree=1,
|
|
ring_degree=1,
|
|
fsdp_dit=False,
|
|
fsdp_text_encoder=False,
|
|
compile_dit=False,
|
|
weight_dtype=torch.bfloat16,
|
|
savedir_sample="samples"
|
|
):
|
|
# 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,
|
|
fsdp_dit=fsdp_dit, fsdp_text_encoder=fsdp_text_encoder, compile_dit=compile_dit,
|
|
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):
|
|
|
|
@app.post("/videox_fun/infer_forward")
|
|
async def _multi_nodes_infer_forward_api(
|
|
datas: dict,
|
|
):
|
|
try:
|
|
result = await engine.generate(datas)
|
|
return result
|
|
except Exception as e:
|
|
if isinstance(e, HTTPException):
|
|
raise e
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
else:
|
|
MultiNodesEngine = None
|
|
MultiNodesGenerator = None
|
|
multi_nodes_infer_forward_api = None |