Add import for SanaVanillaAttnProcessor

Author: scxue

examples/research_projects/sana/train_sana_sprint_diffusers.py

@@ -21,7 +21,7 @@
 import os
 import shutil
 from pathlib import Path
-from typing import Callable
+from typing import Callable, Optional
 
 import accelerate
 import numpy as np
@@ -51,6 +51,7 @@
     SanaSprintPipeline,
     SanaTransformer2DModel,
 )
+from diffusers.models.attention_processor import Attention
 from diffusers.optimization import get_scheduler
 from diffusers.training_utils import (
     free_memory,