Create app.py

app.py

@@ -0,0 +1,125 @@
+import spaces
+import gradio as gr
+import numpy as np
+import random
+import functools
+import os
+import torch
+from diffusers import FluxPipeline
+from peft import LoraConfig, get_peft_model, PeftModel
+
+huggingface_token = os.getenv("HF_TOKEN")
+
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", 
+                                    torch_dtype=torch.float16,
+                                    token=huggingface_token,
+                                    custom_pipeline='quickjkee/swd_pipeline_flux').to('cuda')
+distill_check = 'yresearch/swd_flux'
+pipe.transformer = PeftModel.from_pretrained(
+    pipe.transformer,
+    distill_check,
+)
+
+
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1024
+
+
+@spaces.GPU()
+def infer(prompt, seed, randomize_seed):
+
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+
+    generator = torch.Generator().manual_seed(seed)
+    sigmas = [1.0000, 0.8956, 0.7363, 0.6007, 0.0000]
+    scales = [64, 80, 96, 128]
+    
+    image = pipe(
+        prompt=prompt,
+        height=int(scales[0] * 8),
+        width=int(scales[0] * 8),
+        scales=scales,
+        sigmas=sigmas,
+        timesteps=torch.tensor(sigmas[:-1]).to('cuda') * 1000,
+        guidance_scale=4.5,
+        max_sequence_length=512,
+        generator=torch.Generator("cpu").manual_seed(0)
+    ).images[0]
+
+    return image
+
+
+examples = [
+    "3d digital art of an adorable ghost, holding a heart shaped pumpkin, Halloween, super cute, spooky haunted house background",
+    'Long-exposure night photography of a starry sky over a mountain range, with light trails.',
+    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
+    "A gold astronaut meditating in a lush green forest by a lake",
+    "A group of friends sitting around a campfire."
+]
+
+css = """
+#col-container {
+    margin: 0 auto;
+    max-width: 520px;
+}
+"""
+
+if torch.cuda.is_available():
+    power_device = "GPU"
+else:
+    power_device = "CPU"
+
+with gr.Blocks(css=css) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(
+            f"""
+        # ⚡ Scale-wise Distillation ⚡ 
+        # ⚡ Image Generation with 6-step SwD ⚡
+        This is a demo of [Scale-wise Distillation](https://yandex-research.github.io/swd/), 
+        a diffusion distillation method proposed in [Scale-wise Distillation of Diffusion Models](https://arxiv.org/abs/2503.16397)
+        by [Yandex Research](https://github.com/yandex-research).
+        Currently running on {power_device}.
+        """
+        )
+        gr.Markdown(
+            "If you enjoy the space, feel free to give a ⭐ to the <a href='https://github.com/yandex-research/swd' target='_blank'>Github Repo</a>. [![GitHub Stars](https://img.shields.io/github/stars/yandex-research/invertible-cd?style=social)](https://github.com/yandex-research/invertible-cd)"
+        )
+
+        with gr.Row():
+            prompt = gr.Text(
+                label="Prompt",
+                show_label=False,
+                max_lines=1,
+                placeholder="Enter your prompt",
+                container=False,
+            )
+
+            run_button = gr.Button("Run", scale=0)
+
+        result = gr.Image(label="Result", show_label=False)
+
+        with gr.Accordion("Advanced Settings", open=False):
+            seed = gr.Slider(
+                label="Seed",
+                minimum=0,
+                maximum=MAX_SEED,
+                step=1,
+                value=0,
+            )
+
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
+
+
+        gr.Examples(
+            examples=examples,
+            inputs=[prompt],
+            cache_examples=False
+        )
+    run_button.click(
+        fn=infer,
+        inputs=[prompt, seed, randomize_seed],
+        outputs=[result]
+    )
+
+demo.queue().launch(share=False)