Update app.py

app.py

@@ -1,26 +1,102 @@
-ideo(output.frames[0], path, fps=10)
-    return path
+import gradio as gr
+import torch
+import os
+import spaces
+import uuid
+
+from diffusers import AnimateDiffPipeline, 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
+from transformers import CLIPFeatureExtractor
+
+# 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
+if not torch.cuda.is_available():
+    raise NotImplementedError("No GPU detected!")
+
+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")
+
+feature_extractor = CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32")
+
+@spaces.GPU(duration=30, queue=False)
+def generate_image(prompt, base="Realistic", motion="", step=8, resolution="Square", progress=gr.Progress()):
+    global step_loaded
+    global base_loaded
+    global motion_loaded
+
+    print(prompt, base, step, resolution)
+
+    # Set resolution
+    if resolution == "Square":
+        width, height = 512, 512
+    elif resolution == "Horizontal":
+        width, height = 1280, 720
+    else:
+        width, height = 512, 512  # default fallback
+
+    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))
+
+    output = pipe(
+        prompt=prompt,
+        guidance_scale=1.2,
+        num_inference_steps=step,
+        width=width,
+        height=height,
+        callback=progress_callback,
+        callback_steps=1
+    )
+
+    name = str(uuid.uuid4()).replace("-", "")
+    path = f"/tmp/{name}.mp4"
+    export_to_video(output.frames[0], path, fps=10)
+    return path
 
 # Gradio Interface
 with gr.Blocks(css="style.css") as demo:
-    gr.HTML(
-        "<h1><center>Textual Imagination : A Text To Video Synthesis</center></h1>"
-    )
+    gr.HTML("<h1><center>Textual Imagination : A Text To Video Synthesis</center></h1>")
+    
     with gr.Group():
         with gr.Row():
-            prompt = gr.Textbox(
-                label='Prompt'
-            )
+            prompt = gr.Textbox(label='Prompt')
+        
         with gr.Row():
             select_base = gr.Dropdown(
                 label='Base model',
-                choices=[
-                    "Cartoon", 
-                    "Realistic",
-                    "3d",
-                    "Anime",
-                ],
+                choices=["Cartoon", "Realistic", "3d", "Anime"],
                 value=base_loaded,
                 interactive=True
             )
@@ -53,17 +129,12 @@ with gr.Blocks(css="style.css") as demo:
             )
             select_resolution = gr.Dropdown(
                 label='Resolution',
-                choices=[
-                    "Square", 
-                    "Horizontal",
-                ],
+                choices=["Square", "Horizontal"],
                 value="Square",
                 interactive=True
             )
-            submit = gr.Button(
-                scale=1,
-                variant='primary'
-            )
+            submit = gr.Button(scale=1, variant='primary')
+
     video = gr.Video(
         label='AnimateDiff-Lightning',
         autoplay=True,
@@ -73,10 +144,7 @@ with gr.Blocks(css="style.css") as demo:
     )
 
     gr.on(
-        triggers=[
-            submit.click,
-            prompt.submit
-        ],
+        triggers=[submit.click, prompt.submit],
         fn=generate_image,
         inputs=[prompt, select_base, select_motion, select_step, select_resolution],
         outputs=[video],