import spaces from huggingface_hub import snapshot_download, hf_hub_download import os import subprocess import importlib, site from PIL import Image import uuid import shutil # Re-discover all .pth/.egg-link files for sitedir in site.getsitepackages(): site.addsitedir(sitedir) # Clear caches so importlib will pick up new modules importlib.invalidate_caches() def sh(cmd): subprocess.check_call(cmd, shell=True) flash_attention_installed = False try: print("Attempting to download and install FlashAttention wheel...") flash_attention_wheel = hf_hub_download( repo_id="rahul7star/flash-attn-3", repo_type="model", filename="128/flash_attn_3-3.0.0b1-cp39-abi3-linux_x86_64.whl", ) sh(f"pip install {flash_attention_wheel}") # tell Python to re-scan site-packages now that the egg-link exists import importlib, site; site.addsitedir(site.getsitepackages()[0]); importlib.invalidate_caches() flash_attention_installed = True print("FlashAttention installed successfully.") except Exception as e: print(f"⚠️ Could not install FlashAttention: {e}") print("Continuing without FlashAttention...") import torch print(f"Torch version: {torch.__version__}") print(f"FlashAttention available: {flash_attention_installed}") os.environ["PROCESSED_RESULTS"] = f"{os.getcwd()}/processed_results" import gradio as gr import argparse from ovi.ovi_fusion_engine import OviFusionEngine, DEFAULT_CONFIG from diffusers import DiffusionPipeline import tempfile from ovi.utils.io_utils import save_video from ovi.utils.processing_utils import clean_text, scale_hw_to_area_divisible # ---------------------------- # Parse CLI Args # ---------------------------- parser = argparse.ArgumentParser(description="Ovi Joint Video + Audio Gradio Demo") parser.add_argument( "--cpu_offload", action="store_true", help="Enable CPU offload for both OviFusionEngine and FluxPipeline" ) args = parser.parse_args() ckpt_dir = "./ckpts" # Wan2.2 wan_dir = os.path.join(ckpt_dir, "Wan2.2-TI2V-5B") snapshot_download( repo_id="Wan-AI/Wan2.2-TI2V-5B", local_dir=wan_dir, allow_patterns=[ "google/*", "models_t5_umt5-xxl-enc-bf16.pth", "Wan2.2_VAE.pth" ] ) # MMAudio mm_audio_dir = os.path.join(ckpt_dir, "MMAudio") snapshot_download( repo_id="hkchengrex/MMAudio", local_dir=mm_audio_dir, allow_patterns=[ "ext_weights/best_netG.pt", "ext_weights/v1-16.pth" ] ) ovi_dir = os.path.join(ckpt_dir, "Ovi") snapshot_download( repo_id="chetwinlow1/Ovi", local_dir=ovi_dir, allow_patterns=[ "model.safetensors" ] ) # Initialize OviFusionEngine enable_cpu_offload = args.cpu_offload print(f"loading model...") DEFAULT_CONFIG['cpu_offload'] = enable_cpu_offload # always use cpu offload if image generation is enabled DEFAULT_CONFIG['mode'] = "t2v" # hardcoded since it is always cpu offloaded ovi_engine = OviFusionEngine() try: flux_model = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-Krea-dev", torch_dtype=torch.bfloat16) image_example = None except Exception as e: flux_model = None image_example = "example_prompts/pngs/8.png" print("loaded model") def resize_for_model(image_path): # Open image img = Image.open(image_path) w, h = img.size aspect_ratio = w / h # Decide target size based on aspect ratio if aspect_ratio > 1.5: # wide image target_size = (992, 512) elif aspect_ratio < 0.66: # tall image target_size = (512, 992) else: # roughly square target_size = (512, 512) # Resize while preserving aspect ratio, then pad img.thumbnail(target_size, Image.Resampling.LANCZOS) # Create a new image with target size and paste centered new_img = Image.new("RGB", target_size, (0, 0, 0)) new_img.paste( img, ((target_size[0] - img.size[0]) // 2, (target_size[1] - img.size[1]) // 2) ) return new_img, target_size def _ensure_output_dir(session_id): output_dir = os.path.join(os.environ["PROCESSED_RESULTS"], session_id) os.makedirs(output_dir, exist_ok=True) return output_dir @spaces.GPU() def generate_image(text_prompt, session_id, image_height = 1024, image_width = 1024): """ Generates an image from text_prompt using flux_model if available. Always returns a filepath (string) or raises a gr.Error on failure. """ print("image generation used") text_prompt = clean_text(text_prompt or "") print(text_prompt) # If flux_model isn't loaded, fall back to example image (if available) output_dir = _ensure_output_dir(session_id) output_path = os.path.join(output_dir, "generate_image.png") if flux_model is None: # fallback to example image if provided if image_example and os.path.exists(image_example): # copy example into session folder so downstream can always rely on a path under processed_results shutil.copy(image_example, output_path) print(f"Flux model not available — using example image {image_example}") return output_path else: raise gr.Error("Image generation model not available and no example image found.") # ensure requested dims are divisible/compatible image_h, image_w = scale_hw_to_area_divisible(int(image_height), int(image_width), area=1024 * 1024) try: # move model to GPU, generate, then move model back to CPU flux_model.to("cuda") gen = flux_model( text_prompt, height=image_h, width=image_w, num_inference_steps = 28, guidance_scale=4.5, generator=torch.Generator(device="cuda").manual_seed(1234) ) image = gen.images[0] image.save(output_path) print(f"Saved generated image to {output_path}") return output_path except Exception as e: # provide helpful error message and fallback to example if present print(f"⚠️ generate_image failed: {e}") if image_example and os.path.exists(image_example): shutil.copy(image_example, output_path) print(f"Falling back to example image {image_example}") return output_path raise gr.Error(f"Image generation failed: {e}") finally: try: flux_model.to("cpu") except Exception: pass def generate_scene( text_prompt, sample_steps = 50, image = None, session_id = None, video_seed = 100, solver_name = "unipc", shift = 5, video_guidance_scale = 4, audio_guidance_scale = 3, slg_layer = 11, video_negative_prompt = "", audio_negative_prompt = "", progress=gr.Progress(track_tqdm=True) ): """ Top-level helper that ensures there's an image (generates one if necessary) and then calls generate_video. """ text_prompt_processed = (text_prompt or "").strip() if session_id is None: session_id = uuid.uuid4().hex if not text_prompt_processed: raise gr.Error("Please enter a prompt.") print(text_prompt) # If user did not supply an image (None or empty), try to generate one and use it. if not image: print("No image provided; attempting to generate one.") image = generate_image(text_prompt_processed, session_id) print(f"Generated/fallback image path: {image}") # If image is a dict-like from Gradio, try to extract file path (defensive) if isinstance(image, dict) and "name" in image: image = image["name"] # final check - ensure file exists if not image or not os.path.exists(image): raise gr.Error("No usable image available (generation failed and no fallback).") print(f"{session_id} is generating scene with {sample_steps} steps (image: {image})") return generate_video( text_prompt=text_prompt_processed, sample_steps=sample_steps, image=image, session_id=session_id, video_seed=video_seed, solver_name=solver_name, shift=shift, video_guidance_scale=video_guidance_scale, audio_guidance_scale=audio_guidance_scale, slg_layer=slg_layer, video_negative_prompt=video_negative_prompt, audio_negative_prompt=audio_negative_prompt, progress=progress ) def get_duration( text_prompt, sample_steps, image, session_id, video_seed, solver_name, shift, video_guidance_scale, audio_guidance_scale, slg_layer, video_negative_prompt, audio_negative_prompt, progress, ): image_generation_s = 0 if not image: image_generation_s = 40 warmup = 20 return int(sample_steps * 3 + warmup + image_generation_s) @spaces.GPU(duration=get_duration) def generate_video( text_prompt, sample_steps = 50, image = None, session_id = None, video_seed = 100, solver_name = "unipc", shift = 5, video_guidance_scale = 4, audio_guidance_scale = 3, slg_layer = 11, video_negative_prompt = "", audio_negative_prompt = "", progress=gr.Progress(track_tqdm=True) ): """ Generates a video using ovi_engine given a guaranteed image path (string). """ print("generate_video called") if session_id is None: session_id = uuid.uuid4().hex # If image is not provided for any reason, try generating one now. if not image: print("No image passed to generate_video; generating now...") image = generate_image(text_prompt, session_id) # If Gradio passed a dict or other structure, extract file path if isinstance(image, dict) and "name" in image: image_path = image["name"] else: image_path = image if not image_path or not os.path.exists(image_path): raise gr.Error("Image path is missing or the file does not exist. Cannot generate video.") output_dir = _ensure_output_dir(session_id) output_path = os.path.join(output_dir, "generated_video.mp4") # Resize/pad and get the target dims for the model try: _, target_size = resize_for_model(image_path) except Exception as e: raise gr.Error(f"Failed to open/resize image: {e}") video_frame_width = target_size[0] video_frame_height = target_size[1] # Call your ovi_engine (unchanged) generated_video, generated_audio, _ = ovi_engine.generate( text_prompt=text_prompt, image_path=image_path, video_frame_height_width=[video_frame_height, video_frame_width], seed=video_seed, solver_name=solver_name, sample_steps=sample_steps, shift=shift, video_guidance_scale=video_guidance_scale, audio_guidance_scale=audio_guidance_scale, slg_layer=slg_layer, video_negative_prompt=video_negative_prompt, audio_negative_prompt=audio_negative_prompt, ) save_video(output_path, generated_video, generated_audio, fps=24, sample_rate=16000) print(f"{session_id} video generation succeeded: {output_path}") return output_path def cleanup(request: gr.Request): sid = request.session_hash if sid: d1 = os.path.join(os.environ["PROCESSED_RESULTS"], sid) shutil.rmtree(d1, ignore_errors=True) def start_session(request: gr.Request): return request.session_hash css = """ #col-container { margin: 0 auto; max-width: 1024px; } """ theme = gr.themes.Ocean() with gr.Blocks(css=css, theme=theme) as demo: session_state = gr.State() demo.load(start_session, outputs=[session_state]) with gr.Column(elem_id="col-container"): with gr.Row(): with gr.Column(): # Image section video_text_prompt = gr.Textbox(label="Scene Prompt", lines=5, placeholder="Describe your scene...") sample_steps = gr.Slider( value=20, label="Sample Steps", minimum=20, maximum=100, step=1.0 ) run_btn = gr.Button("Action 🎬", variant="primary") image = gr.Image(type="filepath", label="Image Ref", height=360) gr.Markdown( """ 💡 **Prompt Guidelines** ``` Describe the Scene and Character(s) performance , Dialogue line character voice & atmosphere of the scene ``` """, elem_classes="guideline-bubble" ) with gr.Accordion("🎬 Video Generation Options", open=False, visible=True): video_height = gr.Number(minimum=128, maximum=1280, value=512, step=32, label="Video Height") video_width = gr.Number(minimum=128, maximum=1280, value=992, step=32, label="Video Width") video_seed = gr.Number(minimum=0, maximum=100000, value=100, label="Video Seed") solver_name = gr.Dropdown( choices=["unipc", "euler", "dpm++"], value="unipc", label="Solver Name" ) shift = gr.Slider(minimum=0.0, maximum=20.0, value=5.0, step=1.0, label="Shift") video_guidance_scale = gr.Slider(minimum=0.0, maximum=10.0, value=4.0, step=0.5, label="Video Guidance Scale") audio_guidance_scale = gr.Slider(minimum=0.0, maximum=10.0, value=3.0, step=0.5, label="Audio Guidance Scale") slg_layer = gr.Number(minimum=-1, maximum=30, value=11, step=1, label="SLG Layer") video_negative_prompt = gr.Textbox(label="Video Negative Prompt", placeholder="Things to avoid in video") audio_negative_prompt = gr.Textbox(label="Audio Negative Prompt", placeholder="Things to avoid in audio") with gr.Column(): output_path = gr.Video(label="Generated Video", height=360) gr.Examples( examples=[ [ "What's the difference between having a job and having no life?", 20, "example_prompts/pngs/91.png", ], [ "a alien creature looking to the right and slowly turning to the camera while drooling from her teeth and says Hiss, You thought I can't talk. then start screaming in a high pitch voice the alien has a raspy voice", 20, "example_prompts/pngs/90.png", ], [ "The video opens with a close-up of a woman with vibrant reddish-orange, shoulder-length hair and heavy dark eye makeup. She is wearing a dark brown leather jacket over a grey hooded top. She looks intently to her right, her mouth slightly agape, and her expression is serious and focused. The background shows a room with light green walls and dark wooden cabinets on the left, and a green plant on the right. She speaks, her voice clear and direct, saying, doing. She then pauses briefly, her gaze unwavering, and continues, And I need you to trust them.. Her mouth remains slightly open, indicating she is either about to speak more or has just finished a sentence, with a look of intense sincerity.. Tense, dramatic background music, clear female voice.", 20, image_example, ], [ "A young woman with long, wavy blonde hair and light-colored eyes is shown in a medium shot against a blurred backdrop of lush green foliage. She wears a denim jacket over a striped top. Initially, her eyes are closed and her mouth is slightly open as she speaks, Enjoy this moment. Her eyes then slowly open, looking slightly upwards and to the right, as her expression shifts to one of thoughtful contemplation. She continues to speak, No matter where it's taking, her gaze then settling with a serious and focused look towards someone off-screen to her right.. Clear female voice, faint ambient outdoor sounds.", 20, "example_prompts/pngs/2.png", ], [ "A bearded man wearing large dark sunglasses and a blue patterned cardigan sits in a studio, actively speaking into a large, suspended microphone. He has headphones on and gestures with his hands, displaying rings on his fingers. Behind him, a wall is covered with red, textured sound-dampening foam on the left, and a white banner on the right features the ""CHOICE FM"" logo and various social media handles like ""@ilovechoicefm"" with ""RALEIGH"" below it. The man intently addresses the microphone, articulating, is talent. It's all about authenticity. You gotta be who you really are, especially if you're working. He leans forward slightly as he speaks, maintaining a serious expression behind his sunglasses.. Clear male voice speaking into a microphone, a low background hum.", 20, "example_prompts/pngs/5.png", ], [ "The scene is set outdoors with a blurry, bright green background, suggesting grass and a sunny environment. On the left, a woman with long, dark hair, wearing a red top and a necklace with a white pendant, faces towards the right. Her expression is serious and slightly perturbed as she speaks, with her lips slightly pursed. She says, UFO, UFC thing. On the right, the back of a man's head and his right ear are visible, indicating he is facing away from the camera, listening to the woman. He has short, dark hair. The woman continues speaking, her expression remaining serious, And if you're not watching that, it's one of those ancient movies from an era that's as the frame holds steady on the two figures.. Clear female speech, distant low-frequency hum.", 20, "example_prompts/pngs/9.png", ], ], inputs=[video_text_prompt, sample_steps, image], outputs=[output_path], fn=generate_video, cache_examples=True, ) run_btn.click( fn=generate_scene, inputs=[video_text_prompt, sample_steps, image, session_state], outputs=[output_path], ) if __name__ == "__main__": demo.unload(cleanup) demo.queue() demo.launch(ssr_mode=False, share=True)