Update app.py

app.py

@@ -1,60 +1,63 @@
-import os
-import streamlit as st
 import torch
-from diffusers.utils import load_image
-
-try:
-    from diffusers import CogVideoXImageToVideoPipeline
-    pipeline_available = True
-except ImportError:
-    pipeline_available = False
-    st.error("Failed to import `CogVideoXImageToVideoPipeline`. Please run `pip install diffusers`.")
-
-st.title("Image to Video with Hugging Face")
-st.write("Upload an image and provide a prompt to generate a video.")
-
-if pipeline_available:
-    uploaded_file = st.file_uploader("Upload an image (JPG or PNG):", type=["jpg", "jpeg", "png"])
-    prompt = st.text_input("Enter your prompt:", "A little girl is riding a bicycle at high speed. Focused, detailed, realistic.")
-
-    if uploaded_file and prompt:
-        try:
-            # Save uploaded file
-            import uuid
-            file_name = f"{uuid.uuid4()}_uploaded_image.jpg"
-            with open(file_name, "wb") as f:
-                f.write(uploaded_file.read())
-            st.write("Uploaded image saved successfully.")
-
-            # Load the image
-            image = load_image(file_name)
-
-            # Initialize pipeline
-            device = "cuda" if torch.cuda.is_available() else "cpu"
-            pipe = CogVideoXImageToVideoPipeline.from_pretrained(
-                "THUDM/CogVideoX1.5-5B-I2V",
-                torch_dtype=torch.bfloat16,
-                cache_dir="./huggingface_cache",
-            )
-            pipe.enable_sequential_cpu_offload()
-            pipe.vae.enable_tiling()
-            pipe.vae.enable_slicing()
-
-
-            # Generate video
-            with st.spinner("Generating video... This may take a while."):
-                try:
-                # Attempt to generate the video
-                video_frames = pipe(
-                    prompt=prompt,
-                    image=image,
-                    num_videos_per_prompt=1,
-                    num_inference_steps=50,
-                    num_frames=81,
-                    guidance_scale=6,
-                    generator=torch.Generator(device=device).manual_seed(42),
-                ).frames[0]
-            except Exception as e:
-                # Handle errors gracefully
-                st.error(f"An error occurred during video generation: {e}")
-
+import gradio as gr
+import time
+from diffusers import CogVideoXImageToVideoPipeline
+from diffusers.utils import export_to_video, load_image
+
+# Load model once
+pipe = CogVideoXImageToVideoPipeline.from_pretrained(
+    "THUDM/CogVideoX1.5-5B-I2V",
+    torch_dtype=torch.bfloat16
+)
+pipe.enable_sequential_cpu_offload()
+pipe.vae.enable_tiling()
+pipe.vae.enable_slicing()
+
+def generate_video(image, prompt):
+    if image is None:
+        raise gr.Error("Please upload an input image")
+    if not prompt:
+        raise gr.Error("Please enter a text prompt")
+    
+    # Load uploaded image
+    input_image = load_image(image)
+    
+    # Generate video
+    video_frames = pipe(
+        prompt=prompt,
+        image=input_image,
+        num_videos_per_prompt=1,
+        num_inference_steps=50,
+        num_frames=81,
+        guidance_scale=6,
+        generator=torch.Generator(device="cuda").manual_seed(42),
+    ).frames[0]
+    
+    # Save to temporary file
+    output_path = f"output_{int(time.time())}.mp4"
+    export_to_video(video_frames, output_path, fps=8)
+    
+    return output_path
+
+with gr.Blocks(title="CogVideoX Image-to-Video") as demo:
+    gr.Markdown("# 🎥 CogVideoX Image-to-Video Generation")
+    gr.Markdown("Transform images into videos using AI! Upload an image and enter a description to generate a video.")
+    
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(label="Input Image", type="filepath")
+            prompt_input = gr.Textbox(label="Prompt", placeholder="Describe the video you want to generate...")
+            submit_btn = gr.Button("Generate Video")
+        
+        with gr.Column():
+            video_output = gr.Video(label="Generated Video")
+    gr.Examples(examples=examples, inputs=[image_input, prompt_input])
+    
+    submit_btn.click(
+        fn=generate_video,
+        inputs=[image_input, prompt_input],
+        outputs=video_output,
+    )
+
+if __name__ == "__main__":
+    demo.launch()