fffiloni's picture
add progress bar
afd8ffa verified
raw
history blame
4.54 kB
import gradio as gr
import os
from gradio_client import Client, handle_file
import numpy as np
import tempfile
import imageio
import torch
from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
pipe = DiffusionPipeline.from_pretrained("cerspense/zeroscope_v2_576w", torch_dtype=torch.float16)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe.enable_model_cpu_offload()
hf_token = os.environ.get("HF_TOKEN")
def get_caption(image_in):
kosmos2_client = Client("fffiloni/Kosmos-2-API", hf_token=hf_token)
kosmos2_result = kosmos2_client.predict(
image_input=handle_file(image_in),
text_input="Detailed",
api_name="/generate_predictions"
)
print(f"KOSMOS2 RETURNS: {kosmos2_result}")
data = kosmos2_result[1]
# Extract and combine tokens starting from the second element
sentence = ''.join(item['token'] for item in data[1:])
# Find the last occurrence of "."
#last_period_index = full_sentence.rfind('.')
# Truncate the string up to the last period
#truncated_caption = full_sentence[:last_period_index + 1]
# print(truncated_caption)
#print(f"\n—\nIMAGE CAPTION: {truncated_caption}")
return sentence
def export_to_video(frames: np.ndarray, fps: int) -> str:
frames = np.clip((frames * 255), 0, 255).astype(np.uint8)
out_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
writer = imageio.get_writer(out_file.name, format="FFMPEG", fps=fps)
for frame in frames:
writer.append_data(frame)
writer.close()
return out_file.name
def infer(image_init, progress=gr.Progress(track_tqdm=True)):
prompt = get_caption(image_init)
video_frames = pipe(prompt, num_inference_steps=40, height=320, width=576, num_frames=24).frames[0]
video_path = export_to_video(video_frames, 12)
print(video_path)
return prompt, video_path
css = """
#col-container {max-width: 510px; margin-left: auto; margin-right: auto;}
a {text-decoration-line: underline; font-weight: 600;}
.animate-spin {
animation: spin 1s linear infinite;
}
@keyframes spin {
from {
transform: rotate(0deg);
}
to {
transform: rotate(360deg);
}
}
#share-btn-container {
display: flex;
padding-left: 0.5rem !important;
padding-right: 0.5rem !important;
background-color: #000000;
justify-content: center;
align-items: center;
border-radius: 9999px !important;
max-width: 13rem;
}
#share-btn-container:hover {
background-color: #060606;
}
#share-btn {
all: initial;
color: #ffffff;
font-weight: 600;
cursor:pointer;
font-family: 'IBM Plex Sans', sans-serif;
margin-left: 0.5rem !important;
padding-top: 0.5rem !important;
padding-bottom: 0.5rem !important;
right:0;
}
#share-btn * {
all: unset;
}
#share-btn-container div:nth-child(-n+2){
width: auto !important;
min-height: 0px !important;
}
#share-btn-container .wrap {
display: none !important;
}
#share-btn-container.hidden {
display: none!important;
}
img[src*='#center'] {
display: block;
margin: auto;
}
"""
with gr.Blocks(css=css) as demo:
with gr.Column(elem_id="col-container"):
gr.Markdown(
"""
<h1 style="text-align: center;">Zeroscope Image-to-Video</h1>
<p style="text-align: center;">
A watermark-free Modelscope-based video model optimized for producing high-quality 16:9 compositions and a smooth video output. <br />
This demo is a variation that lets you upload an image as reference for video generation.
</p>
[![Duplicate this Space](https://huggingface.co/datasets/huggingface/badges/raw/main/duplicate-this-space-sm.svg#center)](https://huggingface.co/spaces/fffiloni/zeroscope-img-to-video?duplicate=true)
"""
)
image_init = gr.Image(label="Image Init", type="filepath", sources=["upload"], elem_id="image-init")
#inference_steps = gr.Slider(label="Inference Steps", minimum=10, maximum=100, step=1, value=40, interactive=False)
submit_btn = gr.Button("Submit")
coca_cap = gr.Textbox(label="Caption", placeholder="Kosmos-2 caption will be displayed here", elem_id="coca-cap-in")
video_result = gr.Video(label="Video Output", elem_id="video-output")
submit_btn.click(
fn=infer,
inputs=[image_init],
outputs=[coca_cap, video_result],
show_api=False
)
demo.queue(max_size=12).launch(show_api=False)