update

app.py

@@ -15,7 +15,7 @@ pipe = FluxPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
 pipe = pipe.to(device)
 
 # load pruned model
-pruned_pipe = FluxPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype).to("cpu")
+pruned_pipe = FluxPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
 pruned_pipe.transformer = torch.load(
     hf_hub_download("zhangyang-0123/EcoDiffPrunedModels", "model/flux/flux.pkl"),
     map_location="cpu",
@@ -56,15 +56,18 @@ css = """
 
 header = """
 # 🌱 EcoDiff Pruned FLUX-Schnell (20% Pruning Ratio)
+"""
 
+header_2 = """
 <div style="text-align: center; display: flex; justify-content: left; gap: 5px;">
-<a href="https://arxiv.org/abs/2412.02852"><img src="https://img.shields.io/badge/ariXv-Paper-A42C25.svg" alt="arXiv"></a>
-<a href="https://huggingface.co/zhangyang-0123/EcoDiffPrunedModels"><img src="https://img.shields.io/badge/🤗-Model-ffbd45.svg" alt="HuggingFace"></a>
+    <a href="https://arxiv.org/abs/2412.02852"><img src="https://img.shields.io/badge/ariXv-Paper-A42C25.svg" alt="arXiv"></a>
+    <a href="https://huggingface.co/zhangyang-0123/EcoDiffPrunedModels"><img src="https://img.shields.io/badge/🤗-Model-ffbd45.svg" alt="HuggingFace"></a>
 </div>
 """
 
 with gr.Blocks(css=css) as demo:
     gr.Markdown(header)
+    gr.HTML(header_2)
     with gr.Row():
         prompt = gr.Textbox(
             label="Prompt",