Update app.py

app.py

@@ -1,4 +1,5 @@
 import spaces
+
 """
 Copyright NewGenAI
 Code can't be included in commercial app used for monetary gain. No derivative code allowed.
@@ -11,15 +12,14 @@ import time
 from datetime import datetime
 import os
 
-import torch
 from diffusers.utils import export_to_video
-from diffusers import LTXImageToVideoPipeline
+from diffusers import LTXPipeline
 from transformers import T5EncoderModel, T5Tokenizer
 from pathlib import Path
 from datetime import datetime
 from huggingface_hub import hf_hub_download
 
-STATE_FILE = "LTX091_I2V_state.json"
+STATE_FILE = "LTX091_state.json"
 queue = []
 
 def load_state():
@@ -28,15 +28,16 @@ def load_state():
             return json.load(file)
     return {}
 
+# Function to save the current state
 def save_state(state):
     with open(STATE_FILE, "w") as file:
         json.dump(state, file)
 
+# Load initial state
 initial_state = load_state()
 
-def add_to_queue(image, prompt, negative_prompt, height, width, num_frames, num_inference_steps, fps, seed):
+def add_to_queue(prompt, negative_prompt, height, width, num_frames, num_inference_steps, fps, seed):
     task = {
-        "image": image,
         "prompt": prompt,
         "negative_prompt": negative_prompt,
         "height": height,
@@ -59,7 +60,7 @@ def process_queue():
 
     for i, task in enumerate(queue):
         generate_video(**task)
-        time.sleep(1)
+        time.sleep(1)  # Simulate processing time
 
     queue.clear()
     return "All tasks in the queue have been processed."
@@ -78,7 +79,6 @@ def save_ui_state(prompt, negative_prompt, height, width, num_frames, num_infere
     save_state(state)
     return "State saved!"
 
-# [Previous model loading code remains the same...]
 repo_id = "a-r-r-o-w/LTX-Video-0.9.1-diffusers"
 base_path = repo_id
 files_to_download = [
@@ -102,9 +102,11 @@ files_to_download = [
 os.makedirs(base_path, exist_ok=True)
 for file_path in files_to_download:
     try:
+        # Create the full directory path for this file
         full_dir = os.path.join(base_path, os.path.dirname(file_path))
         os.makedirs(full_dir, exist_ok=True)
         
+        # Download the file
         downloaded_path = hf_hub_download(
             repo_id=repo_id,
             filename=file_path,
@@ -117,20 +119,23 @@ for file_path in files_to_download:
         print(f"Error downloading {file_path}: {str(e)}")
         raise
 
+# Download model from different repo
 try:
+    # Create the full directory path for this file
     full_dir = os.path.join(base_path, os.path.dirname(file_path))
     os.makedirs(full_dir, exist_ok=True)
     
+    # Download the file
     downloaded_path = hf_hub_download(
         repo_id="Lightricks/LTX-Video",
         filename="ltx-video-2b-v0.9.1.safetensors",
         local_dir=repo_id,
     )
-    print(f"Successfully downloaded: ltx-video-2b-v0.9.1.safetensors")
 except Exception as e:
     print(f"Error downloading 0.9.1 model: {str(e)}")
     raise
 
+
 single_file_url = repo_id+"/ltx-video-2b-v0.9.1.safetensors"
 text_encoder = T5EncoderModel.from_pretrained(
   repo_id, subfolder="text_encoder", torch_dtype=torch.bfloat16
@@ -138,7 +143,7 @@ text_encoder = T5EncoderModel.from_pretrained(
 tokenizer = T5Tokenizer.from_pretrained(
   repo_id, subfolder="tokenizer", torch_dtype=torch.bfloat16
 )
-pipe = LTXImageToVideoPipeline.from_single_file(
+pipe = LTXPipeline.from_single_file(
     single_file_url, 
     text_encoder=text_encoder, 
     tokenizer=tokenizer, 
@@ -146,42 +151,42 @@ pipe = LTXImageToVideoPipeline.from_single_file(
 )
 pipe.enable_model_cpu_offload()
 
-@spaces.GPU()
-def generate_video(image, prompt, negative_prompt, height, width, num_frames, num_inference_steps, fps, seed):
+@spaces.GPU(duration=120)
+@torch.inference_mode()
+def generate_video(prompt, negative_prompt, height, width, num_frames, num_inference_steps, fps, seed):
+    # Randomize seed if seed is 0
     if seed == 0:
         seed = random.randint(0, 999999)
-
-    torch.cuda.synchronize()
-    torch.cuda.empty_cache()
-    with torch.inference_mode():
-        video = pipe(
-            image=image,
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            width=width,
-            height=height,
-            num_frames=num_frames,
-            num_inference_steps=num_inference_steps,
-            generator=torch.Generator(device='cuda').manual_seed(seed),
-        ).frames[0]
-
+    
+    # Generating the video <Does not support seed :( >
+    video = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        width=width,
+        height=height,
+        num_frames=num_frames,
+        num_inference_steps=num_inference_steps,
+        generator=torch.Generator(device='cuda').manual_seed(seed),
+    ).frames[0]
+
+    # Create output filename based on prompt and timestamp
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
     filename = f"{prompt[:10]}_{timestamp}.mp4"
     
-    os.makedirs("output_LTX091_i2v", exist_ok=True)
-    output_path = f"./output_LTX091_i2v/{filename}"
+    # Save the video to the output folder
+    os.makedirs("output_LTX091", exist_ok=True)
+    output_path = f"./output_LTX091/{filename}"
     export_to_video(video, output_path, fps=fps)
     
     return output_path
 
+# Gradio UI setup
 def randomize_seed():
     return random.randint(0, 999999)
 
 with gr.Blocks() as demo:
     with gr.Tabs():
         with gr.Tab("Generate Video"):
-            with gr.Row():
-                input_image = gr.Image(label="Input Image", type="pil")
             with gr.Row():
                 prompt = gr.Textbox(label="Prompt", lines=3, value=initial_state.get("prompt", "A dramatic view of the pyramids at Giza during sunset."))
                 negative_prompt = gr.Textbox(label="Negative Prompt", lines=3, value=initial_state.get("negative_prompt", "worst quality, blurry, distorted"))
@@ -203,7 +208,7 @@ with gr.Blocks() as demo:
             random_seed_button.click(lambda: random.randint(0, 999999), outputs=seed)
             generate_button.click(
                 generate_video,
-                inputs=[input_image, prompt, negative_prompt, height, width, num_frames, num_inference_steps, fps, seed],
+                inputs=[prompt, negative_prompt, height, width, num_frames, num_inference_steps, fps, seed],
                 outputs=output_video
             )
             save_state_button.click(
@@ -213,8 +218,6 @@ with gr.Blocks() as demo:
             )
 
         with gr.Tab("Batch Processing"):
-            with gr.Row():
-                batch_input_image = gr.Image(label="Input Image", type="pil")
             with gr.Row():
                 batch_prompt = gr.Textbox(label="Prompt", lines=3, value="A batch of videos depicting different landscapes.")
                 batch_negative_prompt = gr.Textbox(label="Negative Prompt", lines=3, value="low quality, inconsistent, jittery")
@@ -238,7 +241,7 @@ with gr.Blocks() as demo:
             random_seed_batch_button.click(lambda: random.randint(0, 999999), outputs=batch_seed)
             add_to_queue_button.click(
                 add_to_queue,
-                inputs=[batch_input_image, batch_prompt, batch_negative_prompt, batch_height, batch_width, batch_num_frames, batch_num_inference_steps, batch_fps, batch_seed],
+                inputs=[batch_prompt, batch_negative_prompt, batch_height, batch_width, batch_num_frames, batch_num_inference_steps, batch_fps, batch_seed],
                 outputs=queue_status
             )
             clear_queue_button.click(clear_queue, outputs=queue_status)