more cleanup

app.py

@@ -4,6 +4,8 @@ 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():
@@ -40,6 +42,8 @@ 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
@@ -52,11 +56,7 @@ 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(
-    "--use_image_gen",
-    action="store_true",
-    help="Enable image generation UI with FluxPipeline"
-)
+
 parser.add_argument(
     "--cpu_offload",
     action="store_true",
@@ -99,16 +99,11 @@ snapshot_download(
 )
 
 # Initialize OviFusionEngine
-enable_cpu_offload = args.cpu_offload or args.use_image_gen
-use_image_gen = args.use_image_gen
-print(f"loading model... {enable_cpu_offload=}, {use_image_gen=} for gradio demo")
+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()
-flux_model = None
-if use_image_gen:
-    flux_model = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-Krea-dev", torch_dtype=torch.bfloat16)
-    flux_model.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU VRAM
 print("loaded model")
 
 
@@ -170,6 +165,7 @@ def generate_video(
     slg_layer = 11,
     video_negative_prompt = "",
     audio_negative_prompt = "",
+    session_id = None,
     progress=gr.Progress(track_tqdm=True)
 ):
     try:
@@ -178,6 +174,15 @@ def generate_video(
         if image is not None:
             image_path = image
 
+        if session_id is None:
+            session_id = uuid.uuid4().hex
+            
+
+        output_dir = os.path.join(os.environ["PROCESSED_RESULTS"], session_id)
+        os.makedirs(output_dir, exist_ok=True)
+        output_path = os.path.join(output_dir, f"generated_video.mp4")
+
+
         _, target_size = resize_for_model(image_path)
 
         video_frame_width = target_size[0]
@@ -198,8 +203,6 @@ def generate_video(
             audio_negative_prompt=audio_negative_prompt,
         )
 
-        tmpfile = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
-        output_path = tmpfile.name
         save_video(output_path, generated_video, generated_audio, fps=24, sample_rate=16000)
 
         return output_path
@@ -208,24 +211,16 @@ def generate_video(
         return None
 
 
-def generate_image(text_prompt, image_seed, image_height, image_width):
-    if flux_model is None:
-        return None
-    text_prompt = clean_text(text_prompt)
-    print(f"Generating image with prompt='{text_prompt}', seed={image_seed}, size=({image_height},{image_width})")
-
-    image_h, image_w = scale_hw_to_area_divisible(image_height, image_width, area=1024 * 1024)
-    image = flux_model(
-        text_prompt,
-        height=image_h,
-        width=image_w,
-        guidance_scale=4.5,
-        generator=torch.Generator().manual_seed(int(image_seed))
-    ).images[0]
-
-    tmpfile = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
-    image.save(tmpfile.name)
-    return tmpfile.name
+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 {
@@ -236,6 +231,9 @@ css = """
 
 with gr.Blocks(css=css) as demo:
 
+    session_state = gr.State()
+    demo.load(start_session, outputs=[session_state])
+
     with gr.Column(elem_id="col-container"):
         gr.HTML(
             """
@@ -337,13 +335,6 @@ with gr.Blocks(css=css) as demo:
                     cache_examples=True,
                 )
 
-    if args.use_image_gen and gen_img_btn is not None:
-        gen_img_btn.click(
-            fn=generate_image,
-            inputs=[image_text_prompt, image_seed, image_height, image_width],
-            outputs=[image],
-        )
-
     run_btn.click(
         fn=generate_video,
         inputs=[video_text_prompt, image, sample_steps],
@@ -351,4 +342,7 @@ with gr.Blocks(css=css) as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.unload(cleanup)
+    demo.queue()
+    demo.launch(ssr_mode=False, share=True)
+