Update app.py

app.py

@@ -1,92 +1,32 @@
-import os
-
-# Set CUDA environment variables _before_ any torch imports.
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"
-os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
-
+import streamlit as st
 import subprocess
-import sys
-import torch
-
-# For debugging, print CUDA availability.
-print("torch.cuda.is_available():", torch.cuda.is_available())
-
-# If torch reports no CUDA devices, apply monkey-patches.
-if not torch.cuda.is_available():
-    print("Monkey patching torch.cuda to force CUDA availability.")
-    torch.cuda.is_available = lambda: True
-    torch.cuda.current_device = lambda: 0
-    torch.cuda.set_device = lambda device: None
-    # Override lazy initialization to do nothing.
-    torch.cuda._lazy_init = lambda: None
-
-# Try to force initialization (will now use our patch instead of raising an error).
-try:
-    _ = torch.cuda.current_device()
-    print("Current CUDA device (patched):", torch.cuda.current_device())
-except Exception as e:
-    print("Error forcing CUDA initialization:", e)
-
-# List of required packages.
-required_packages = ["easydict", "diffusers", "ftfy", "transformers"]
-
-for package in required_packages:
-    try:
-        __import__(package)
-    except ModuleNotFoundError:
-        print(f"{package} not found, installing now...")
-        subprocess.check_call([sys.executable, "-m", "pip", "install", package])
-
-from easydict import EasyDict
-import gradio as gr
-from huggingface_hub import snapshot_download
-
-# Define local directory names.
-REPO_DIR = "Wan2_1"  # Renamed from Wan2.1 to avoid invalid module names.
-CHECKPOINT_PARENT_DIR = "model_checkpoints"
-CHECKPOINT_SUBDIR = "Wan2.1-T2V-14B"  # Example checkpoint subdirectory.
-
-# Clone the repository if it does not exist.
-if not os.path.exists(REPO_DIR):
-    print("Cloning Wan2.1 repository...")
-    subprocess.run(["git", "clone", "https://github.com/Wan-Video/Wan2.1.git", REPO_DIR])
-
-# Add the cloned repository to Python's module search path.
-sys.path.insert(0, os.path.abspath(REPO_DIR))
-
-# Download the model checkpoint snapshot from Hugging Face Hub.
-os.makedirs(CHECKPOINT_PARENT_DIR, exist_ok=True)
-checkpoint_dir = os.path.join(CHECKPOINT_PARENT_DIR, CHECKPOINT_SUBDIR)
-print("Downloading model checkpoint from Hugging Face Hub...")
-model_id = "Wan-AI/Wan2.1-T2V-14B"  # Update if necessary.
-ckpt_path = snapshot_download(repo_id=model_id, cache_dir=checkpoint_dir)
-print(f"Model checkpoint downloaded to: {ckpt_path}")
-
-# Import WanPipeline from the local repository.
-from wan.pipeline import WanPipeline
-
-# Initialize the pipeline.
-pipe = WanPipeline(
-    model_dir=ckpt_path,  # Directory containing model files.
-    device="cuda" if torch.cuda.is_available() else "cpu"
-)
-
-def generate_video(prompt):
-    """
-    Generate a video from a text prompt using Wan2.1.
-    Calls the pipeline's generate() method and returns the video frames.
-    """
-    video_frames = pipe.generate(prompt)
-    return video_frames
-
-# Create the Gradio interface.
-iface = gr.Interface(
-    fn=generate_video,
-    inputs=gr.Textbox(label="Prompt", placeholder="Enter a video prompt here..."),
-    outputs=gr.Video(label="Generated Video"),
-    title="Wan2.1 Text-to-Video Generation",
-    description="Generate videos from text prompts using the Wan2.1 model."
-)
+import os
 
-if __name__ == "__main__":
-    iface.launch()
+# Title
+st.title("🎥 WAN 2.1 - 14B AI Text-to-Video Generator")
+
+# Input fields
+prompt = st.text_area("Enter your text prompt:", "A cat in military dress wearing headphones, laughing and walking.")
+frame_num = st.slider("Number of frames:", min_value=30, max_value=120, value=60, step=10)
+resolution = st.selectbox("Select resolution:", ["832*480", "1280*720"])
+sample_steps = st.slider("Sampling steps:", min_value=10, max_value=50, value=20, step=5)
+
+# Button to generate video
+if st.button("Generate Video"):
+    st.info("Generating video... This may take a few minutes.")
+
+    # Run WAN 2.1 - 14B Model
+    command = f"python generate.py --task t2v-14B --size {resolution} --frame_num {frame_num} --sample_steps {sample_steps} --ckpt_dir ./Wan2.1-T2V-14B --offload_model True --prompt \"{prompt}\""
+    
+    process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    stdout, stderr = process.communicate()
+
+    # Print logs for debugging
+    st.text_area("📜 Logs", stdout.decode() + stderr.decode())
+
+    # Check if video was created
+    if os.path.exists("output.mp4"):
+        st.video("output.mp4")
+        st.success("✅ Video generated successfully!")
+    else:
+        st.error("❌ Video generation failed! Check logs above.")