CrazyEric/flash-sd3

1import random2import spaces3 4import gradio as gr5import numpy as np6import torch7from diffusers import StableDiffusion3Pipeline, SD3Transformer2DModel, FlashFlowMatchEulerDiscreteScheduler8from peft import PeftModel9import os10from huggingface_hub import snapshot_download11 12huggingface_token = os.getenv("HUGGINFACE_TOKEN")13 14model_path = snapshot_download(15    repo_id="stabilityai/stable-diffusion-3-medium", 16    revision="refs/pr/26",17    repo_type="model", 18    ignore_patterns=["*.md", "*..gitattributes"],19    local_dir="stable-diffusion-3-medium",20    token=huggingface_token, # type a new token-id.21    )22 23device = "cuda" if torch.cuda.is_available() else "cpu"24IS_SPACE = os.environ.get("SPACE_ID", None) is not None25 26transformer = SD3Transformer2DModel.from_pretrained(27    model_path,28    subfolder="transformer",29    torch_dtype=torch.float16,30)31transformer = PeftModel.from_pretrained(transformer, "jasperai/flash-sd3")32 33 34if torch.cuda.is_available():35    torch.cuda.max_memory_allocated(device=device)36    pipe = StableDiffusion3Pipeline.from_pretrained(37        model_path,38        transformer=transformer,39        torch_dtype=torch.float16,40        text_encoder_3=None,41        tokenizer_3=None,42    )43 44    pipe = pipe.to(device)45else:46    pipe = StableDiffusion3Pipeline.from_pretrained(47        model_path,48        transformer=transformer,49        torch_dtype=torch.float16,50        text_encoder_3=None,51        tokenizer_3=None,52    )53    pipe = pipe.to(device)54 55 56pipe.scheduler = FlashFlowMatchEulerDiscreteScheduler.from_pretrained(57  model_path,58  subfolder="scheduler",59)60 61MAX_SEED = np.iinfo(np.int32).max62MAX_IMAGE_SIZE = 102463NUM_INFERENCE_STEPS = 464 65 66@spaces.GPU67def infer(prompt, seed, randomize_seed, guidance_scale, num_inference_steps, negative_prompt, progress=gr.Progress(track_tqdm=True)):68    if randomize_seed:69        seed = random.randint(0, MAX_SEED)70 71    generator = torch.Generator().manual_seed(seed)72 73    image = pipe(74        prompt=prompt,75        guidance_scale=guidance_scale,76        num_inference_steps=num_inference_steps,77        generator=generator,78        negative_prompt=negative_prompt79    ).images[0]80 81    return image82 83 84examples = [85    "The image showcases a freshly baked bread, possibly focaccia, with rosemary sprigs and red pepper flakes sprinkled on top. It's sliced and placed on a wire cooling rack, with a bowl of mixed peppercorns beside it.",86    'a 3D render of a wizard raccoon holding a sign saying "SD 3" with a magic wand.',87    "A panda reading a book in a lush forest.",88    "A raccoon trapped inside a glass jar full of colorful candies, the background is steamy with vivid colors",89    "Pirate ship sailing on a sea with the milky way galaxy in the sky and purple glow lights",90    "a cute cartoon fluffy rabbit pilot walking on a military aircraft carrier, 8k, cinematic",91    "A 3d render of a futuristic city with a giant robot in the middle full of neon lights, pink and blue colors",92    "A close up of an old elderly man with green eyes looking straight at the camera",93    "photo of a huge red cat with green eyes sitting on a cloud in the sky, looking at the camera"94]95 96css = """97#col-container {98    margin: 0 auto;99    max-width: 512px;100}101"""102 103if torch.cuda.is_available():104    power_device = "GPU"105else:106    power_device = "CPU"107 108with gr.Blocks(css=css) as demo:109    with gr.Column(elem_id="col-container"):110        gr.Markdown(111            f"""112        # ⚡ Flash Diffusion: FlashSD3 ⚡113        This is an interactive demo of [Flash Diffusion](https://gojasper.github.io/flash-diffusion-project/), a diffusion distillation method proposed in [Flash Diffusion: Accelerating Any Conditional114        Diffusion Model for Few Steps Image Generation](http://arxiv.org/abs/2406.02347) *by Clément Chadebec, Onur Tasar, Eyal Benaroche and Benjamin Aubin* from Jasper Research.115        [This model](https://huggingface.co/jasperai/flash-sd3) is a **90.4M** LoRA distilled version of [SD3](https://huggingface.co/stabilityai/stable-diffusion-3-medium) model that is able to generate 1024x1024 images in **4 to 8 steps**.116        Results can be compared with the teacher model [here](https://huggingface.co/spaces/stabilityai/stable-diffusion-3-medium).117        Currently running on {power_device}.118        """119        )120        gr.Markdown(121            "If you enjoy the space, please also promote *open-source* by giving a ⭐ to the <a href='https://github.com/gojasper/flash-diffusion' target='_blank'>Github Repo</a>. [![GitHub Stars](https://img.shields.io/github/stars/gojasper/flash-diffusion?style=social)](https://github.com/gojasper/flash-diffusion)"122        )123        gr.Markdown(124            "💡 *Hint:* We noticed that 8 steps and CFG can improve the results (text rendering in particular) for that very model. Feel free to play with those parameters."125        )126 127        gr.Markdown(128            "💡 *Hint:* To better appreciate the low latency of our method, run the demo locally !"129        )130 131        with gr.Row():132            prompt = gr.Text(133                label="Prompt",134                show_label=False,135                max_lines=1,136                placeholder="Enter your prompt",137                container=False,138            )139 140            run_button = gr.Button("Run", scale=0)141 142        result = gr.Image(label="Result", show_label=False)143 144        with gr.Accordion("Advanced Settings", open=False):145 146            negative_prompt = gr.Text(147                label="Negative prompt",148                max_lines=1,149                placeholder="Enter a negative prompt",150                value="deformed, distorted, disfigured, poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, mutated hands and fingers, disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation, NSFW, bad text"151            )152            153            seed = gr.Slider(154                label="Seed",155                minimum=0,156                maximum=MAX_SEED,157                step=1,158                value=0,159            )160 161            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)162 163            with gr.Row():164                165                guidance_scale = gr.Slider(166                    label="Guidance scale",167                    minimum=0.0,168                    maximum=3.0,169                    step=0.1,170                    value=1.0,171                )172                173                num_inference_steps = gr.Slider(174                    label="Number of inference steps",175                    minimum=4,176                    maximum=8,177                    step=1,178                    value=4,179                )180 181        examples = gr.Examples(examples=examples, inputs=[prompt], cache_examples=False)182 183        gr.Markdown("**Disclaimer:**")184        gr.Markdown(185            "This demo is only for research purpose. Jasper cannot be held responsible for the generation of NSFW (Not Safe For Work) content through the use of this demo. Users are solely responsible for any content they create, and it is their obligation to ensure that it adheres to appropriate and ethical standards. Jasper provides the tools, but the responsibility for their use lies with the individual user."186        )187    gr.on(188        [run_button.click, seed.change, randomize_seed.change, prompt.submit],189        fn=infer,190        inputs=[prompt, seed, randomize_seed, guidance_scale, num_inference_steps, negative_prompt],191        outputs=[result],192        # show_progress="minimal",193        #show_api=False,194        #trigger_mode="always_last",195    )196 197demo.queue().launch(show_api=False)198