Update app.py

app.py

@@ -1,42 +1,40 @@
 import gradio as gr
 from diffusers import StableDiffusionPipeline, DiffusionPipeline
-import torch, imageio
+import torch
+from PIL import Image
 
-# Load Image Generator
-img_model = "runwayml/stable-diffusion-v1-5"
-img_pipe = StableDiffusionPipeline.from_pretrained(img_model, torch_dtype=torch.float32)
-img_pipe = img_pipe.to("cpu")
+# Load Text ➤ Image model
+text2img = StableDiffusionPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16
+).to("cuda")
 
-# Load Video Generator
-vid_model = "damo-vilab/text-to-video-ms-1.7b"
-vid_pipe = DiffusionPipeline.from_pretrained(vid_model, torch_dtype=torch.float32)
-vid_pipe = vid_pipe.to("cpu")
+# Load Image ➤ Video model
+img2vid = DiffusionPipeline.from_pretrained(
+    "damo-vilab/image-to-video", torch_dtype=torch.float16, variant="fp16"
+).to("cuda")
 
-# Image generation function
 def generate_image(prompt):
-    image = img_pipe(prompt).images[0]
+    image = text2img(prompt).images[0]
     return image
 
-# Video generation function
-def generate_video(prompt):
-    video_frames = vid_pipe(prompt, num_frames=8).frames
-    output_path = "output.mp4"
-    imageio.mimsave(output_path, video_frames, fps=8)
-    return output_path
+def generate_video(image):
+    video_frames = img2vid(image).frames
+    return video_frames
 
-# Gradio Tabs UI
 with gr.Blocks() as demo:
-    gr.Markdown("## 🎨 Free AI Image + Video Generator")
-    with gr.Tab("🖼️ Image Generator"):
-        text_in = gr.Textbox(label="ছবির জন্য প্রম্পট লিখুন")
-        img_out = gr.Image(type="pil")
-        btn1 = gr.Button("Generate Image")
-        btn1.click(generate_image, inputs=text_in, outputs=img_out)
-
-    with gr.Tab("🎬 Video Generator"):
-        text_in2 = gr.Textbox(label="ভিডিওর জন্য প্রম্পট লিখুন")
-        vid_out = gr.Video()
-        btn2 = gr.Button("Generate Video")
-        btn2.click(generate_video, inputs=text_in2, outputs=vid_out)
+    gr.Markdown("## 🎨 Text ➤ Image ➤ Video Generator")
+
+    with gr.Row():
+        prompt = gr.Textbox(label="📝 Enter Prompt")
+        img_output = gr.Image(label="🖼️ Generated Image")
+        btn_img = gr.Button("Generate Image")
+
+    with gr.Row():
+        img_input = gr.Image(label="📥 Drop Image Here")
+        vid_output = gr.Video(label="🎬 Generated Video")
+        btn_vid = gr.Button("Generate Video")
+
+    btn_img.click(fn=generate_image, inputs=prompt, outputs=img_output)
+    btn_vid.click(fn=generate_video, inputs=img_input, outputs=vid_output)
 
 demo.launch()