Update app.py

app.py

@@ -57,14 +57,14 @@ def generate_image(prompt, base="Realistic", motion="", step=1, progress=gr.Prog
     logger.info(f"Generating video with prompt: {prompt}, base: {base}, steps: {step}")
 
     try:
-        # Tải checkpoint AnimateDiff-Lightning
+        # Tải AnimateDiff Lightning checkpoint
         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
 
-        # Tải mô hình cơ sở
+        # Tải mô hình cơ sở nếu thay đổi
         if base_loaded != base:
             pipe.unet.load_state_dict(
                 torch.load(hf_hub_download(bases[base], "unet/diffusion_pytorch_model.bin"), map_location=device),
@@ -72,66 +72,59 @@ def generate_image(prompt, base="Realistic", motion="", step=1, progress=gr.Prog
             )
             base_loaded = base
 
-        # Tải motion LoRA (tùy chọn)
+        # Tải motion LoRA nếu có
         if motion_loaded != motion:
             try:
                 pipe.unload_lora_weights()
-                if motion != "":
+                if motion:
                     pipe.load_lora_weights(motion, adapter_name="motion")
                     pipe.set_adapters(["motion"], [0.7])
                 motion_loaded = motion
             except Exception as e:
-                logger.warning(f"Failed to load LoRA weights: {str(e)}. Proceeding without LoRA.")
+                logger.warning(f"Failed to load motion LoRA: {e}")
                 motion_loaded = ""
 
         progress((0, step))
         def progress_callback(i, t, z):
             progress((i + 1, step))
 
-        # Tối ưu hóa suy luận
+        # Suy luận
         with torch.no_grad():
             output = pipe(
                 prompt=prompt,
                 guidance_scale=1.2,
                 num_inference_steps=step,
-                num_frames=32,  # Tạo 32 khung hình, sẽ lặp lại để đạt 192
+                num_frames=32,
                 callback=progress_callback,
                 callback_steps=1,
                 width=256,
                 height=256
             )
 
-        # Chuẩn bị khung hình
-        frames = output.frames[0]  # frames là list các numpy arrays
-        target_fps = 24
-        target_duration = 8  # 8 giây
-        target_frames = target_fps * target_duration  # 192 khung hình
-
-        # Lặp lại khung hình để đạt 192
-        current_frames = len(frames)
-        if current_frames < target_frames:
-            repeat_factor = (target_frames + current_frames - 1) // current_frames
-            frames = np.repeat(frames, repeat_factor, axis=0)[:target_frames]
-        elif current_frames > target_frames:
+        # Chuẩn hóa khung hình cho 8 giây
+        frames = output.frames[0]
+        fps = 24
+        target_frames = fps * 8
+        if len(frames) < target_frames:
+            frames = np.tile(frames, (target_frames // len(frames) + 1, 1, 1, 1))[:target_frames]
+        else:
             frames = frames[:target_frames]
 
-        logger.info(f"Generated {len(frames)} frames for video.")
-
-        # Xuất video với imageio
+        # Tạo video
         name = str(uuid.uuid4()).replace("-", "")
-        path = os.path.join(output_dir, f"{name}.mp4")
-        logger.info(f"Saving video to {path}")
-        export_to_video(frames, path, fps=target_fps)
+        video_path = os.path.join(output_dir, f"{name}.mp4")
+        export_to_video(frames, video_path, fps=fps)
+
+        if not os.path.exists(video_path):
+            raise FileNotFoundError("❌ Video was not created")
+
+        logger.info(f"✅ Video ready at {video_path}")
 
-        # Kiểm tra file video
-        if not os.path.exists(path):
-            raise FileNotFoundError(f"Video file {path} was not created.")
-        
-        logger.info(f"Video generated successfully: {path}")
-        return path
+        # ✅ Trả về đối tượng `gr.File` để có link tải được từ client/frontend
+        return gr.File(video_path, file_name=f"{name}.mp4", show_download_button=True)
 
     except Exception as e:
-        logger.error(f"Error generating video: {str(e)}")
+        logger.error(f"❌ Error in generation: {e}")
         raise
 
 # Giao diện Gradio