Spaces:
Runtime error
Runtime error
| 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 <https://github.com/lopho/makeavid-sd-tpu> | |
| """) | |
| 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)) | |
| <small>Leave a ❤️ like if you like. Consider it a dopamine donation at no cost.</small> | |
| """) | |
| 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() | |