fix app

app.py

@@ -1,13 +1,13 @@
 import gradio as gr
 import numpy as np
-import random
+from huggingface_hub import hf_hub_download
 
 import spaces  # [uncomment to use ZeroGPU]
 from diffusers import DiffusionPipeline
 import torch
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
+model_repo_id = "stabilityai/stable-diffusion-xl-base-1.0"  # Replace to the model you would like to use
 
 if torch.cuda.is_available():
     torch_dtype = torch.float16
@@ -17,6 +17,15 @@ else:
 pipe = DiffusionPipeline.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.transformer = torch.load(
+    hf_hub_download("zhangyang-0123/EcoDiffPrunedModels", "model/sdxl/sdxl.pkl"),
+    map_location="cpu",
+    torch_dtype=torch_dtype,
+)
+pruned_pipe = pruned_pipe.to(device)
+
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024