date banner

README.md

@@ -9,8 +9,6 @@ app_file: app.py
 pinned: true
 python_version: 3.10
 short_description: Diffusion Model Compression
-hf_oauth: true
-hf_oauth_expiration_minutes: 43200
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -38,19 +38,15 @@ def binary_mask_eval(args, model):
             "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16
         ).to("cpu")
         pruned_pipe.unet = torch.load(
-            hf_hub_download(
-                "zhangyang-0123/EcoDiffPrunedModels", "model/sdxl/sdxl.pkl"
-            ),
+            hf_hub_download("zhangyang-0123/EcoDiffPrunedModels", "model/sdxl/sdxl.pkl"),
             map_location="cpu",
         )
     elif model == "flux":
-        pruned_pipe = FluxPipeline.from_pretrained(
-            "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16
-        ).to("cpu")
+        pruned_pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to(
+            "cpu"
+        )
         pruned_pipe.transformer = torch.load(
-            hf_hub_download(
-                "zhangyang-0123/EcoDiffPrunedModels", "model/flux/flux.pkl"
-            ),
+            hf_hub_download("zhangyang-0123/EcoDiffPrunedModels", "model/flux/flux.pkl"),
             map_location="cpu",
         )
 
@@ -60,9 +56,7 @@ def binary_mask_eval(args, model):
             "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16
         ).to("cpu")
     elif model == "flux":
-        pipe = FluxPipeline.from_pretrained(
-            "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16
-        ).to("cpu")
+        pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to("cpu")
 
     print("prune complete")
     return pipe, pruned_pipe
@@ -74,13 +68,9 @@ def generate_images(prompt, seed, steps, pipe, pruned_pipe):
     pruned_pipe.to("cuda")
     # Run the model and return images directly
     g_cpu = torch.Generator("cuda").manual_seed(seed)
-    original_image = pipe(
-        prompt=prompt, generator=g_cpu, num_inference_steps=steps
-    ).images[0]
+    original_image = pipe(prompt=prompt, generator=g_cpu, num_inference_steps=steps).images[0]
     g_cpu = torch.Generator("cuda").manual_seed(seed)
-    ecodiff_image = pruned_pipe(
-        prompt=prompt, generator=g_cpu, num_inference_steps=steps
-    ).images[0]
+    ecodiff_image = pruned_pipe(prompt=prompt, generator=g_cpu, num_inference_steps=steps).images[0]
     return original_image, ecodiff_image
 
 
@@ -91,15 +81,29 @@ def on_prune_click(prompt, seed, steps, model):
 
 
 def on_generate_click(prompt, seed, steps, pipe, pruned_pipe):
-    original_image, ecodiff_image = generate_images(
-        prompt, seed, steps, pipe, pruned_pipe
-    )
+    original_image, ecodiff_image = generate_images(prompt, seed, steps, pipe, pruned_pipe)
     return original_image, ecodiff_image
 
 
+header = """
+# 🌱 Text-to-Image Generation with EcoDiff Pruned Models
+<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://github.com/YaNgZhAnG-V5/EcoDiff"><img src="https://img.shields.io/badge/GitHub-Code-blue.svg?logo=github&" alt="GitHub"></a>
+</div>
+
+<div style="text-align: center; display: flex; justify-content: left; gap: 5px;">
+For **⚡faster⚡** DEMO on one model only, please visit
+<a href="https://huggingface.co/spaces/zhangyang-0123/EcoDiff-SD-XL"><img alt="Static Badge" src="https://img.shields.io/badge/SDXL-fedcba.svg"></a>
+<a href="https://huggingface.co/spaces/zhangyang-0123/EcoDiff-FLUX-Schnell"><img alt="Static Badge" src="https://img.shields.io/badge/FLUX-fgdfba"></a>
+</div>
+"""
+
+
 def create_demo():
     with gr.Blocks() as demo:
-        gr.Markdown("# Text-to-Image Generation with EcoDiff Pruned Model")
+        gr.Markdown(header)
         with gr.Row():
             gr.Markdown(
                 """
@@ -107,18 +111,10 @@ def create_demo():
                 """
             )
         with gr.Row():
-            model_choice = gr.Radio(
-                choices=["SDXL", "FLUX"], value="SDXL", label="Model", scale=2
-            )
-            pruning_ratio = gr.Text(
-                "20% Pruning Ratio for SDXL, FLUX", label="Pruning Ratio", scale=2
-            )
-            status_label = gr.HighlightedText(
-                label="Model Status", value=[("Model Not Initialized", "red")], scale=1
-            )
-            prune_btn = gr.Button(
-                "Initialize Original and Pruned Models", variant="primary", scale=1
-            )
+            model_choice = gr.Radio(choices=["SDXL", "FLUX"], value="SDXL", label="Model", scale=2)
+            pruning_ratio = gr.Text("20% Pruning Ratio for SDXL, FLUX", label="Pruning Ratio", scale=2)
+            status_label = gr.HighlightedText(label="Model Status", value=[("Model Not Initialized", "red")], scale=1)
+            prune_btn = gr.Button("Initialize Original and Pruned Models", variant="primary", scale=1)
 
         with gr.Row():
             gr.Markdown(
@@ -182,4 +178,4 @@ def create_demo():
 
 if __name__ == "__main__":
     demo = create_demo()
-    demo.launch(share=True)
+    demo.launch(share=True)