Update app.py

app.py

@@ -1,62 +1,26 @@
-import streamlit as st
-from diffusers import StableDiffusionPipeline
-from PIL import Image
-import torch
-import io
-import base64
-
-# Model loading (cached)
-@st.cache_resource
-def load_model():
-    pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to("cuda" if torch.cuda.is_available() else "cpu")
-    return pipe
-
-pipe = load_model()
-
-st.title("AI Image Generator")
-
-prompt = st.text_area("Enter your prompt:", height=150)
-num_images = st.slider("Number of Images", 1, 4, 1)
-
-if st.button("Generate"):
-    if prompt:
-        with st.spinner("Generating images..."):
-            images = pipe(prompt, num_images=num_images).images
-
-        for i, image in enumerate(images):
-            col1, col2 = st.columns([1, 2])
-
-            with col1:
-                st.image(image, caption=f"Image {i+1}")
-
-            with col2:
-                buffered = io.BytesIO()
-                image.save(buffered, format="PNG")
-                img_str = base64.b64encode(buffered.getvalue()).decode()
-                href = f'<a href="data:image/png;base64,{img_str}" download="image_{i+1}.png">Download Image {i+1}</a>'
-                st.markdown(href, unsafe_allow_html=True)
-
-    else:
-        st.warning("Please enter a prompt.")
-
-# Styling (optional)
-st.markdown("""
-<style>
-body {
-    font-family: sans-serif;
-}
-.stButton>button {
-    background-color: #4CAF50;
-    border: none;
-    color: white;
-    padding: 10px 20px;
-    text-align: center;
-    text-decoration: none;
-    display: inline-block;
-    font-size: 16px;
-    margin: 4px 2px;
-    cursor: pointer;
-    border-radius: 5px;
-}
-</style>
-""", unsafe_allow_html=True)
+import imageio  
+import imageio_ffmpeg  
+import torch  
+from diffusers import MochiPipeline  
+from diffusers.utils import export_to_video  
+  
+# Load the pre-trained video generation model  
+model = MochiPipeline.from_pretrained(  
+    "MISHANM/video_generation",  
+ #   variant="bf16",  
+    torch_dtype=torch.bfloat16,  
+    device_map="balanced"  
+)  
+  
+# Enable memory savings by tiling the VAE  
+model.enable_vae_tiling()  
+  
+# Define the prompt and number of frames  
+prompt = "A cow drinking water on the surface of Mars."  
+num_frames = 20  
+  
+frames = model(prompt, num_frames=num_frames).frames[0]  
+    
+export_to_video(frames, "video.mp4", fps=30)  
+  
+print("Video generation complete. Saved as 'video.mp4'.")