Create main.py

main.py

@@ -0,0 +1,138 @@
+```python
+import os
+import uuid
+import logging
+import torch
+import numpy as np
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+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 fastapi.responses import FileResponse
+
+# Thiết lập logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+app = FastAPI()
+
+# Tạo thư mục lưu video
+output_dir = "/app/outputs"
+os.makedirs(output_dir, exist_ok=True)
+
+# Constants
+bases = {
+    "Cartoon": "frankjoshua/toonyou_beta6",
+    "Realistic": "emilianJR/epiCRealism",
+    "3d": "Lykon/DreamShaper",
+    "Anime": "Yntec/mistoonAnime2"
+}
+step_loaded = None
+base_loaded = "Realistic"
+motion_loaded = None
+
+# Thiết lập thiết bị CPU và kiểu dữ liệu
+device = "cpu"
+dtype = torch.float32
+
+# Khởi tạo pipeline
+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"
+)
+pipe.safety_checker = None
+
+# Mô hình dữ liệu cho request
+class VideoRequest(BaseModel):
+    prompt: str
+    base: str = "Realistic"
+    motion: str = ""
+    step: int = 1
+
+# Endpoint tạo video
+@app.post("/generate_video")
+async def generate_video(request: VideoRequest):
+    global step_loaded, base_loaded, motion_loaded
+    prompt = request.prompt
+    base = request.base
+    motion = request.motion
+    step = request.step
+
+    logger.info(f"Tạo video với prompt: {prompt}, base: {base}, motion: {motion}, steps: {step}")
+
+    try:
+        # Kiểm tra base hợp lệ
+        if base not in bases:
+            raise HTTPException(status_code=400, detail="Base model không hợp lệ")
+
+        # 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ở 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),
+                strict=False
+            )
+            base_loaded = base
+
+        # Tải motion LoRA nếu có
+        if motion_loaded != motion:
+            try:
+                pipe.unload_lora_weights()
+                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"Không thể tải motion LoRA: {e}")
+                motion_loaded = ""
+
+        # Suy luận
+        with torch.no_grad():
+            output = pipe(
+                prompt=prompt,
+                guidance_scale=1.2,
+                num_inference_steps=step,
+                num_frames=32,
+                width=256,
+                height=256
+            )
+
+        # 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]
+
+        # Tạo video
+        name = str(uuid.uuid4()).replace("-", "")
+        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 không được tạo")
+
+        logger.info(f"Video sẵn sàng tại {video_path}")
+
+        # Trả về file video
+        return FileResponse(video_path, media_type="video/mp4", filename=f"{name}.mp4")
+
+    except Exception as e:
+        logger.error(f"Lỗi khi tạo video: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+# Endpoint kiểm tra trạng thái
+@app.get("/")
+async def root():
+    return {"message": "FastAPI AnimateDiff-Lightning API on Hugging Face Spaces"}
+```