import os import json from io import BytesIO import base64 from functools import partial from PIL import Image, ImageOps import gradio as gr from makeavid_sd.inference import ( InferenceUNetPseudo3D, jnp, SCHEDULERS ) print(os.environ.get('XLA_PYTHON_CLIENT_PREALLOCATE', 'NotSet')) print(os.environ.get('XLA_PYTHON_CLIENT_ALLOCATOR', 'NotSet')) _seen_compilations = set() _model = InferenceUNetPseudo3D( model_path = 'TempoFunk/makeavid-sd-jax', dtype = jnp.float16, hf_auth_token = os.environ.get('HUGGING_FACE_HUB_TOKEN', None) ) if _model.failed != False: trace = f'```{_model.failed}```' with gr.Blocks(title = 'Make-A-Video Stable Diffusion JAX', analytics_enabled = False) as demo: exception = gr.Markdown(trace) demo.launch() _examples = [] _expath = 'examples' for x in sorted(os.listdir(_expath)): with open(os.path.join(_expath, x, 'params.json'), 'r') as f: ex = json.load(f) ex['image_input'] = None if os.path.isfile(os.path.join(_expath, x, 'input.png')): ex['image_input'] = os.path.join(_expath, x, 'input.png') ex['image_output'] = os.path.join(_expath, x, 'output.gif') _examples.append(ex) _output_formats = ( 'webp', 'gif' ) # gradio is illiterate. type hints make it go poopoo in pantsu. def generate( prompt = 'An elderly man having a great time in the park.', neg_prompt = '', hint_image = None, inference_steps = 20, cfg = 15.0, cfg_image = 9.0, seed = 0, fps = 12, num_frames = 24, height = 512, width = 512, scheduler_type = 'dpm', output_format = 'gif' ) -> str: num_frames = min(24, max(2, int(num_frames))) inference_steps = min(60, max(2, int(inference_steps))) height = min(576, max(256, int(height))) width = min(576, max(256, int(width))) height = (height // 64) * 64 width = (width // 64) * 64 cfg = max(cfg, 1.0) cfg_image = max(cfg_image, 1.0) fps = min(1000, max(1, int(fps))) seed = min(2**32-2, int(seed)) if seed < 0: seed = -seed if hint_image is not None: if hint_image.mode != 'RGB': hint_image = hint_image.convert('RGB') if hint_image.size != (width, height): hint_image = ImageOps.fit(hint_image, (width, height), method = Image.Resampling.LANCZOS) scheduler_type = scheduler_type.lower() if scheduler_type not in SCHEDULERS: scheduler_type = 'dpm' output_format = output_format.lower() if output_format not in _output_formats: output_format = 'gif' mask_image = None images = _model.generate( prompt = [prompt] * _model.device_count, neg_prompt = neg_prompt, hint_image = hint_image, mask_image = mask_image, inference_steps = inference_steps, cfg = cfg, cfg_image = cfg_image, height = height, width = width, num_frames = num_frames, seed = seed, scheduler_type = scheduler_type ) _seen_compilations.add((hint_image is None, inference_steps, height, width, num_frames)) with BytesIO() as buffer: images[1].save( buffer, format = output_format, save_all = True, append_images = images[2:], loop = 0, duration = round(1000 / fps), allow_mixed = True, optimize = True ) data = f'data:image/{output_format};base64,' + base64.b64encode(buffer.getvalue()).decode() with BytesIO() as buffer: images[-1].save(buffer, format = 'png', optimize = True) last_data = f'data:image/png;base64,' + base64.b64encode(buffer.getvalue()).decode() with BytesIO() as buffer: images[0].save(buffer, format ='png', optimize = True) first_data = f'data:image/png;base64,' + base64.b64encode(buffer.getvalue()).decode() return data, last_data, first_data def check_if_compiled(hint_image, inference_steps, height, width, num_frames, scheduler_type, message): height = int(height) width = int(width) inference_steps = int(inference_steps) height = (height // 64) * 64 width = (width // 64) * 64 if (hint_image is None, inference_steps, height, width, num_frames, scheduler_type) in _seen_compilations: return '' else: return message with gr.Blocks(title = 'Make-A-Video Stable Diffusion JAX', analytics_enabled = False) as demo: variant = 'panel' with gr.Row(): with gr.Column(): intro1 = gr.Markdown(""" # Make-A-Video Stable Diffusion JAX We have extended a pretrained latent-diffusion inpainting image generation model with **temporal convolutions and attention**. We guide the video generation with a hint image by taking advantage of the extra 5 input channels of the inpainting model. In this demo the hint image can be given by the user, otherwise it is generated by an generative image model. The temporal layers are a port of [Make-A-Video PyTorch](https://github.com/lucidrains/make-a-video-pytorch) to [JAX](https://github.com/google/jax) utilizing [FLAX](https://github.com/google/flax). The convolution is pseudo 3D and seperately convolves accross the spatial dimension in 2D and over the temporal dimension in 1D. Temporal attention is purely self attention and also separately attends to time. Only the new temporal layers have been fine tuned on a dataset of videos themed around dance. The model has been trained for 80 epochs on a dataset of 18,000 Videos with 120 frames each, randomly selecting a 24 frame range from each sample. See model and dataset links in the metadata. Model implementation and training code can be found at """) with gr.Column(): intro3 = gr.Markdown(""" **Please be patient. The model might have to compile with current parameters.** This can take up to 5 minutes on the first run, and 2-3 minutes on later runs. The compilation will be cached and later runs with the same parameters will be much faster. Changes to the following parameters require the model to compile - Number of frames - Width & Height - Inference steps - Input image vs. no input image - Noise scheduler type If you encounter any issues, please report them here: [Space discussions](https://huggingface.co/spaces/TempoFunk/makeavid-sd-jax/discussions) (or DM [@lopho](https://twitter.com/lopho)) Leave a ❤️ like if you like. Consider it a dopamine donation at no cost. """) with gr.Row(variant = variant): with gr.Column(): with gr.Row(): #cancel_button = gr.Button(value = 'Cancel') submit_button = gr.Button(value = 'Make A Video', variant = 'primary') prompt_input = gr.Textbox( label = 'Prompt', value = 'They are dancing in the club but everybody is a 3d cg hairy monster wearing a hairy costume.', interactive = True ) neg_prompt_input = gr.Textbox( label = 'Negative prompt (optional)', value = 'monochrome, saturated', interactive = True ) cfg_input = gr.Slider( label = 'Guidance scale video', minimum = 1.0, maximum = 20.0, step = 0.1, value = 15.0, interactive = True ) cfg_image_input = gr.Slider( label = 'Guidance scale hint (no effect with input image)', minimum = 1.0, maximum = 20.0, step = 0.1, value = 15.0, interactive = True ) seed_input = gr.Number( label = 'Random seed', value = 0, interactive = True, precision = 0 ) image_input = gr.Image( label = 'Hint image (optional)', interactive = True, image_mode = 'RGB', type = 'pil', optional = True, source = 'upload' ) inference_steps_input = gr.Slider( label = 'Steps', minimum = 2, maximum = 60, value = 20, step = 1, interactive = True ) num_frames_input = gr.Slider( label = 'Number of frames to generate', minimum = 2, maximum = 24, step = 1, value = 24, interactive = True ) width_input = gr.Slider( label = 'Width', minimum = 256, maximum = 576, step = 64, value = 512, interactive = True ) height_input = gr.Slider( label = 'Height', minimum = 256, maximum = 576, step = 64, value = 512, interactive = True ) scheduler_input = gr.Dropdown( label = 'Noise scheduler', choices = list(SCHEDULERS.keys()), value = 'dpm', interactive = True ) with gr.Row(): fps_input = gr.Slider( label = 'Output FPS', minimum = 1, maximum = 1000, step = 1, value = 12, interactive = True ) output_format = gr.Dropdown( label = 'Output format', choices = _output_formats, value = 'gif', interactive = True ) with gr.Column(): #will_trigger = gr.Markdown('') patience = gr.Markdown('**Please be patient. The model might have to compile with current parameters.**') image_output = gr.Image( label = 'Output', value = 'example.gif', interactive = False ) tips = gr.Markdown('🤫 *Secret tip*: try using the last frame as input for the next generation.') with gr.Row(): last_frame_output = gr.Image( label = 'Last frame', interactive = False ) first_frame_output = gr.Image( label = 'Initial frame', interactive = False ) examples_lst = [] for x in _examples: examples_lst.append([ x['image_output'], x['prompt'], x['neg_prompt'], x['image_input'], x['cfg'], x['cfg_image'], x['seed'], x['fps'], x['steps'], x['scheduler'], x['num_frames'], x['height'], x['width'], x['format'] ]) examples = gr.Examples( examples = examples_lst, inputs = [ image_output, prompt_input, neg_prompt_input, image_input, cfg_input, cfg_image_input, seed_input, fps_input, inference_steps_input, scheduler_input, num_frames_input, height_input, width_input, output_format ], postprocess = False ) #trigger_inputs = [ image_input, inference_steps_input, height_input, width_input, num_frames_input, scheduler_input ] #trigger_check_fun = partial(check_if_compiled, message = 'Current parameters need compilation.') #height_input.change(fn = trigger_check_fun, inputs = trigger_inputs, outputs = will_trigger) #width_input.change(fn = trigger_check_fun, inputs = trigger_inputs, outputs = will_trigger) #num_frames_input.change(fn = trigger_check_fun, inputs = trigger_inputs, outputs = will_trigger) #image_input.change(fn = trigger_check_fun, inputs = trigger_inputs, outputs = will_trigger) #inference_steps_input.change(fn = trigger_check_fun, inputs = trigger_inputs, outputs = will_trigger) #scheduler_input.change(fn = trigger_check_fun, inputs = trigger_inputs, outputs = will_trigger) submit_button.click( fn = generate, inputs = [ prompt_input, neg_prompt_input, image_input, inference_steps_input, cfg_input, cfg_image_input, seed_input, fps_input, num_frames_input, height_input, width_input, scheduler_input, output_format ], outputs = [ image_output, last_frame_output, first_frame_output ], postprocess = False ) #cancel_button.click(fn = lambda: None, cancels = ev) demo.queue(concurrency_count = 1, max_size = 8) demo.launch()