|
|
| 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 Exception:
|
| 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,
|
| ):
|
|
|
| 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"
|
| ):
|
|
|
| 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 |
