Spaces:
Runtime error
Runtime error
from PIL import Image | |
import os | |
import cv2 | |
import numpy as np | |
from PIL import Image | |
from moviepy.editor import * | |
import gradio as gr | |
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler | |
import torch | |
torch.backends.cuda.matmul.allow_tf32 = True | |
import gc | |
controlnet = ControlNetModel.from_pretrained("ioclab/control_v1p_sd15_brightness", torch_dtype=torch.float16, use_safetensors=True) | |
pipe = StableDiffusionControlNetPipeline.from_pretrained( | |
"runwayml/stable-diffusion-v1-5", | |
controlnet=controlnet, | |
torch_dtype=torch.float16, | |
safety_checker=None, | |
) | |
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) | |
pipe.enable_xformers_memory_efficient_attention() | |
pipe.enable_model_cpu_offload() | |
pipe.enable_attention_slicing() | |
def get_frames(video_in): | |
frames = [] | |
#resize the video | |
clip = VideoFileClip(video_in) | |
#check fps | |
if clip.fps > 30: | |
print("vide rate is over 30, resetting to 30") | |
clip_resized = clip.resize(height=512) | |
clip_resized.write_videofile("video_resized.mp4", fps=30) | |
else: | |
print("video rate is OK") | |
clip_resized = clip.resize(height=512) | |
clip_resized.write_videofile("video_resized.mp4", fps=clip.fps) | |
print("video resized to 512 height") | |
# Opens the Video file with CV2 | |
cap= cv2.VideoCapture("video_resized.mp4") | |
fps = cap.get(cv2.CAP_PROP_FPS) | |
print("video fps: " + str(fps)) | |
i=0 | |
while(cap.isOpened()): | |
ret, frame = cap.read() | |
if ret == False: | |
break | |
cv2.imwrite('kang'+str(i)+'.jpg',frame) | |
frames.append('kang'+str(i)+'.jpg') | |
i+=1 | |
cap.release() | |
cv2.destroyAllWindows() | |
print("broke the video into frames") | |
return frames, fps | |
def create_video(frames, fps): | |
print("building video result") | |
clip = ImageSequenceClip(frames, fps=fps) | |
clip.write_videofile("_result.mp4", fps=fps) | |
return "_result.mp4" | |
def process_brightness( | |
prompt, | |
negative_prompt, | |
conditioning_image, | |
num_inference_steps=30, | |
size=768, | |
guidance_scale=7.0, | |
seed=1234, | |
): | |
conditioning_image_raw = Image.fromarray(conditioning_image) | |
conditioning_image = conditioning_image_raw.convert('L') | |
g_cpu = torch.Generator() | |
if seed == -1: | |
generator = g_cpu.manual_seed(g_cpu.seed()) | |
else: | |
generator = g_cpu.manual_seed(seed) | |
output_image = pipe( | |
prompt, | |
conditioning_image, | |
height=size, | |
width=size, | |
num_inference_steps=num_inference_steps, | |
generator=generator, | |
negative_prompt=negative_prompt, | |
guidance_scale=guidance_scale, | |
controlnet_conditioning_scale=1.0, | |
).images[0] | |
del conditioning_image, conditioning_image_raw | |
gc.collect() | |
return output_image | |
def infer(video_in, prompt, | |
negative_prompt, | |
conditioning_image, | |
num_inference_steps=30, | |
size=768, | |
guidance_scale=7.0, | |
seed=1234 | |
): | |
# 1. break video into frames and get FPS | |
break_vid = get_frames(video_in) | |
frames_list= break_vid[0] | |
fps = break_vid[1] | |
#n_frame = int(trim_value*fps) | |
n_frame = len(frames_list) | |
if n_frame >= len(frames_list): | |
print("video is shorter than the cut value") | |
n_frame = len(frames_list) | |
# 2. prepare frames result arrays | |
result_frames = [] | |
print("set stop frames to: " + str(n_frame)) | |
for i, image in enumerate(frames_list[0:int(n_frame)]): | |
image = Image.open(image).convert("RGB") | |
image = np.array(image) | |
output_frame = process_brightness(image, | |
prompt, | |
negative_prompt, | |
conditioning_image, | |
num_inference_steps=30, | |
size=768, | |
guidance_scale=7.0, | |
seed=1234 | |
) | |
print(output_frame) | |
image = Image.open(output_frame) | |
#image = Image.fromarray(output_frame[0]) | |
image.save("_frame_" + str(i) + ".jpeg") | |
result_frames.append("_frame_" + str(i) + ".jpeg") | |
print("frame " + str(i) + "/" + str(n_frame) + ": done;") | |
final_vid = create_video(result_frames, fps) | |
return final_vid | |
with gr.Blocks() as demo: | |
gr.Markdown( | |
""" | |
# ControlNet on Brightness • Video | |
This is a demo on ControlNet based on brightness for video. | |
""") | |
with gr.Row(): | |
with gr.Column(): | |
prompt = gr.Textbox( | |
label="Prompt", | |
) | |
negative_prompt = gr.Textbox( | |
label="Negative Prompt", | |
) | |
video_in = gr.Video( | |
label="Conditioning Video", | |
source="upload", | |
type="filepath" | |
) | |
with gr.Accordion('Advanced options', open=False): | |
with gr.Row(): | |
num_inference_steps = gr.Slider( | |
10, 40, 20, | |
step=1, | |
label="Steps", | |
) | |
size = gr.Slider( | |
256, 768, 512, | |
step=128, | |
label="Size", | |
) | |
with gr.Row(): | |
guidance_scale = gr.Slider( | |
label='Guidance Scale', | |
minimum=0.1, | |
maximum=30.0, | |
value=7.0, | |
step=0.1 | |
) | |
seed = gr.Slider( | |
label='Seed', | |
value=-1, | |
minimum=-1, | |
maximum=2147483647, | |
step=1, | |
# randomize=True | |
) | |
submit_btn = gr.Button( | |
value="Submit", | |
variant="primary" | |
) | |
with gr.Column(min_width=300): | |
output = gr.Image( | |
label="Result", | |
) | |
submit_btn.click( | |
fn=infer, | |
inputs=[ | |
video_in, prompt, negative_prompt, conditioning_image, num_inference_steps, size, guidance_scale, seed | |
], | |
outputs=output | |
) | |
gr.Examples( | |
examples=[ | |
["a village in the mountains", "monochrome", "./conditioning_images/conditioning_image_1.jpg"], | |
["three people walking in an alleyway with hats and pants", "monochrome", "./conditioning_images/conditioning_image_2.jpg"], | |
["an anime character, natural skin", "monochrome, blue skin, grayscale", "./conditioning_images/conditioning_image_3.jpg"], | |
["a man in a black suit", "monochrome", "./conditioning_images/conditioning_image_4.jpg"], | |
["the forbidden city in beijing at sunset with a reflection in the water", "monochrome", "./conditioning_images/conditioning_image_5.jpg"], | |
["a man in a white shirt holding his hand out in front of", "monochrome", "./conditioning_images/conditioning_image_6.jpg"], | |
], | |
inputs=[ | |
prompt, negative_prompt, conditioning_image | |
], | |
outputs=output, | |
fn=infer, | |
cache_examples=False, | |
) | |
gr.Markdown( | |
""" | |
* [Dataset](https://huggingface.co/datasets/ioclab/grayscale_image_aesthetic_3M) | |
* [Diffusers model](https://huggingface.co/ioclab/control_v1p_sd15_brightness), [Web UI model](https://huggingface.co/ioclab/ioc-controlnet) | |
* [Training Report](https://api.wandb.ai/links/ciaochaos/oot5cui2), [Doc(Chinese)](https://aigc.ioclab.com/sd-showcase/brightness-controlnet.html) | |
""") | |
demo.launch() |