update

app.py

@@ -3,7 +3,7 @@ import numpy as np
 from huggingface_hub import hf_hub_download
 
 import spaces  # [uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
+from diffusers import StableDiffusionXLPipeline
 import torch
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -11,11 +11,11 @@ model_repo_id = "stabilityai/stable-diffusion-xl-base-1.0"  # Replace to the mod
 
 torch_dtype = torch.bfloat16
 
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
+pipe = StableDiffusionXLPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
 pipe = pipe.to(device)
 
 # load pruned model
-pruned_pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
+pruned_pipe = StableDiffusionXLPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
 pruned_pipe.unet = torch.load(
     hf_hub_download("zhangyang-0123/EcoDiffPrunedModels", "model/sdxl/sdxl.pkl"),
     map_location="cpu",