Update app.py

app.py

@@ -1,149 +1,109 @@
 import gradio as gr
 import torch
 import os
-import spaces
 import uuid
 
 from diffusers import AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler
 from diffusers.utils import export_to_video
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
-from PIL import Image
 
-# Constants
-bases = {
-    "Cartoon": "frankjoshua/toonyou_beta6",
-    "Realistic": "emilianJR/epiCRealism", 
-    "3d": "Lykon/DreamShaper",
-    "Anime": "Yntec/mistoonAnime2"
-}
-step_loaded = None
-base_loaded = "Realistic"
-motion_loaded = None
-
-# Ensure model and scheduler are initialized in GPU-enabled function
+# Ensure GPU Availability
 if not torch.cuda.is_available():
-    raise NotImplementedError("No GPU detected!")
+    raise NotImplementedError("A GPU is required for this task.")
 
 device = "cuda"
 dtype = torch.float16
-pipe = AnimateDiffPipeline.from_pretrained(bases[base_loaded], torch_dtype=dtype).to(device)
-pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", beta_schedule="linear")
-
-# Safety checkers
-from transformers import CLIPFeatureExtractor
-
-feature_extractor = CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32")
-
-# Function 
-@spaces.GPU(duration=30,queue=False)
-def generate_image(prompt, base="Realistic", motion="", step=8, progress=gr.Progress()):
-    global step_loaded
-    global base_loaded
-    global motion_loaded
-    print(prompt, base, step)
-
-    if step_loaded != step:
-        repo = "ByteDance/AnimateDiff-Lightning"
-        ckpt = f"animatediff_lightning_{step}step_diffusers.safetensors"
-        pipe.unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device), strict=False)
-        step_loaded = step
-
-    if base_loaded != base:
-        pipe.unet.load_state_dict(torch.load(hf_hub_download(bases[base], "unet/diffusion_pytorch_model.bin"), map_location=device), strict=False)
-        base_loaded = base
-
-    if motion_loaded != motion:
-        pipe.unload_lora_weights()
-        if motion != "":
-            pipe.load_lora_weights(motion, adapter_name="motion")
-            pipe.set_adapters(["motion"], [0.7])
-        motion_loaded = motion
 
-    progress((0, step))
-    def progress_callback(i, t, z):
-        progress((i+1, step))
+# Base Model Paths
+BASE_MODELS = {
+    "Realistic": "emilianJR/epiCRealism",
+    "Cartoon": "frankjoshua/toonyou_beta6",
+    "3D": "Lykon/DreamShaper",
+    "Anime": "Yntec/mistoonAnime2",
+}
 
-    output = pipe(prompt=prompt, guidance_scale=1.2, num_inference_steps=step, callback=progress_callback, callback_steps=1)
+# Initialize Pipeline
+print("Loading AnimateDiff pipeline...")
+base_model = "Realistic"
+pipe = AnimateDiffPipeline.from_pretrained(BASE_MODELS[base_model], torch_dtype=dtype).to(device)
+pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", beta_schedule="linear")
+print("Pipeline loaded successfully.")
+
+# Video Generation Function
+def generate_video(prompt, base="Realistic", motion="", steps=8):
+    global pipe
+    print(f"Generating video: Prompt='{prompt}', Base='{base}', Steps='{steps}'")
+
+    # Switch Base Model
+    if base in BASE_MODELS:
+        print(f"Loading base model: {base}")
+        pipe = AnimateDiffPipeline.from_pretrained(BASE_MODELS[base], torch_dtype=dtype).to(device)
+
+    # Set Inference Steps
+    steps = int(steps)
+    fps = 10  # Frames per second
+    duration = 30  # Video duration in seconds
+    total_frames = fps * duration  # Total frames to generate
+
+    # Generate Frames
+    video_frames = []
+    for i in range(total_frames):
+        output = pipe(
+            prompt=prompt,
+            guidance_scale=1.2,
+            num_inference_steps=steps
+        )
+        video_frames.extend(output.frames[0])
 
+    # Export to Video
     name = str(uuid.uuid4()).replace("-", "")
-    path = f"/tmp/{name}.mp4"
-    export_to_video(output.frames[0], path, fps=10)
-    return path
-
+    output_path = f"/tmp/{name}.mp4"
+    export_to_video(video_frames, output_path, fps=fps)
+    print(f"Video saved to {output_path}")
+    return output_path
 
 # Gradio Interface
-with gr.Blocks(css="style.css") as demo:
-    gr.HTML(
-        "<h1><center>Textual Imagination : A Text To Video Synthesis</center></h1>"
+with gr.Blocks() as demo:
+    gr.HTML("<h1><center>30-Second Text-to-Video Generation</center></h1>")
+    
+    with gr.Row():
+        prompt = gr.Textbox(label="Text Prompt", placeholder="Describe your scene...")
+    
+    with gr.Row():
+        base_model = gr.Dropdown(
+            label="Base Model",
+            choices=["Realistic", "Cartoon", "3D", "Anime"],
+            value="Realistic"
         )
-    with gr.Group():
-        with gr.Row():
-            prompt = gr.Textbox(
-                label='Prompt'
-            )
-        with gr.Row():
-            select_base = gr.Dropdown(
-                label='Base model',
-                choices=[
-                    "Cartoon", 
-                    "Realistic",
-                    "3d",
-                    "Anime",
-                ],
-                value=base_loaded,
-                interactive=True
-            )
-            select_motion = gr.Dropdown(
-                label='Motion',
-                choices=[
-                    ("Default", ""),
-                    ("Zoom in", "guoyww/animatediff-motion-lora-zoom-in"),
-                    ("Zoom out", "guoyww/animatediff-motion-lora-zoom-out"),
-                    ("Tilt up", "guoyww/animatediff-motion-lora-tilt-up"),
-                    ("Tilt down", "guoyww/animatediff-motion-lora-tilt-down"),
-                    ("Pan left", "guoyww/animatediff-motion-lora-pan-left"),
-                    ("Pan right", "guoyww/animatediff-motion-lora-pan-right"),
-                    ("Roll left", "guoyww/animatediff-motion-lora-rolling-anticlockwise"),
-                    ("Roll right", "guoyww/animatediff-motion-lora-rolling-clockwise"),
-                ],
-                value="guoyww/animatediff-motion-lora-zoom-in",
-                interactive=True
-            )
-            select_step = gr.Dropdown(
-                label='Inference steps',
-                choices=[
-                    ('1-Step', 1), 
-                    ('2-Step', 2),
-                    ('4-Step', 4),
-                    ('8-Step', 8),
-                ],
-                value=4,
-                interactive=True
-            )
-            submit = gr.Button(
-                scale=1,
-                variant='primary'
-            )
-    video = gr.Video(
-        label='AnimateDiff-Lightning',
-        autoplay=True,
-        height=512,
-        width=512,
-        elem_id="video_output"
-    )
-
-    gr.on(triggers=[
-            submit.click,
-            prompt.submit
-    ],
-        fn = generate_image,
-        inputs = [prompt, select_base, select_motion, select_step],
-        outputs = [video],
-        api_name = "instant_video",
-        queue = False
+        motion = gr.Dropdown(
+            label="Motion Adapter",
+            choices=[
+                ("None", ""),
+                ("Zoom In", "guoyww/animatediff-motion-lora-zoom-in"),
+                ("Zoom Out", "guoyww/animatediff-motion-lora-zoom-out"),
+                ("Tilt Up", "guoyww/animatediff-motion-lora-tilt-up"),
+                ("Tilt Down", "guoyww/animatediff-motion-lora-tilt-down"),
+                ("Pan Left", "guoyww/animatediff-motion-lora-pan-left"),
+                ("Pan Right", "guoyww/animatediff-motion-lora-pan-right"),
+            ],
+            value=""
+        )
+        steps = gr.Dropdown(
+            label="Inference Steps",
+            choices=["4", "8", "12"],
+            value="8"
+        )
+    
+    with gr.Row():
+        generate_button = gr.Button("Generate Video")
+    
+    video_output = gr.Video(label="Generated Video", autoplay=True, height=512, width=512)
+
+    generate_button.click(
+        fn=generate_video,
+        inputs=[prompt, base_model, motion, steps],
+        outputs=video_output
     )
 
-demo.queue().launch()
-
-Translate
+demo.launch()