Update app.py

app.py

@@ -1,351 +1,149 @@
+import random
 import spaces
+
 import gradio as gr
-from gradio_imageslider import ImageSlider
+import numpy as np
 import torch
-
-torch.jit.script = lambda f: f
-from diffusers import (
-    ControlNetModel,
-    StableDiffusionXLControlNetImg2ImgPipeline,
-    DDIMScheduler,
-)
-from controlnet_aux import AnylineDetector
-from compel import Compel, ReturnedEmbeddingsType
-from PIL import Image
+from diffusers import LCMScheduler, PixArtAlphaPipeline, Transformer2DModel
+from peft import PeftModel
 import os
-import time
-import numpy as np
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
 IS_SPACES_ZERO = os.environ.get("SPACES_ZERO_GPU", "0") == "1"
-IS_SPACE = os.environ.get("SPACE_ID", None) is not None
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-dtype = torch.float16
+transformer = Transformer2DModel.from_pretrained(
+    "PixArt-alpha/PixArt-XL-2-1024-MS",
+    subfolder="transformer",
+    torch_dtype=torch.float16,
+)
+transformer = PeftModel.from_pretrained(transformer, "jasperai/flash-pixart")
 
-LOW_MEMORY = os.getenv("LOW_MEMORY", "0") == "1"
 
-print(f"device: {device}")
-print(f"dtype: {dtype}")
-print(f"low memory: {LOW_MEMORY}")
+if torch.cuda.is_available():
+    torch.cuda.max_memory_allocated(device=device)
+    pipe = PixArtAlphaPipeline.from_pretrained(
+        "PixArt-alpha/PixArt-XL-2-1024-MS",
+        transformer=transformer,
+        torch_dtype=torch.float16,
+    )
+    if not IS_SPACES_ZERO:
+        pipe.enable_xformers_memory_efficient_attention()
+    pipe = pipe.to(device)
+else:
+    pipe = PixArtAlphaPipeline.from_pretrained(
+        "PixArt-alpha/PixArt-XL-2-1024-MS",
+        transformer=transformer,
+        torch_dtype=torch.float16,
+    )
+    pipe = pipe.to(device)
 
+pipe.text_encoder.to_bettertransformer()
 
-model = "stabilityai/stable-diffusion-xl-base-1.0"
-# model = "stabilityai/sdxl-turbo"
-# vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=dtype)
-scheduler = DDIMScheduler.from_pretrained(model, subfolder="scheduler")
-# controlnet = ControlNetModel.from_pretrained(
-#     "diffusers/controlnet-canny-sdxl-1.0", torch_dtype=torch.float16
-# )
-controlnet = ControlNetModel.from_pretrained(
-    "TheMistoAI/MistoLine",
-    torch_dtype=torch.float16,
-    revision="refs/pr/3",
-    variant="fp16",
-)
-pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
-    model,
-    controlnet=controlnet,
-    torch_dtype=dtype,
-    variant="fp16",
-    use_safetensors=True,
-    scheduler=scheduler,
+pipe.scheduler = LCMScheduler.from_pretrained(
+    "PixArt-alpha/PixArt-XL-2-1024-MS",
+    subfolder="scheduler",
+    timestep_spacing="trailing",
 )
 
-compel = Compel(
-    tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
-    text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
-    returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
-    requires_pooled=[False, True],
-)
-pipe = pipe.to(device)
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1024
+NUM_INFERENCE_STEPS = 4
 
-anyline = AnylineDetector.from_pretrained(
-    "TheMistoAI/MistoLine", filename="MTEED.pth", subfolder="Anyline"
-).to(device)
 
+@spaces.GPU
+def infer(prompt, seed, randomize_seed):
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
 
-def pad_image(image):
-    w, h = image.size
-    if w == h:
-        return image
-    elif w > h:
-        new_image = Image.new(image.mode, (w, w), (0, 0, 0))
-        pad_w = 0
-        pad_h = (w - h) // 2
-        new_image.paste(image, (0, pad_h))
-        return new_image
-    else:
-        new_image = Image.new(image.mode, (h, h), (0, 0, 0))
-        pad_w = (h - w) // 2
-        pad_h = 0
-        new_image.paste(image, (pad_w, 0))
-        return new_image
+    generator = torch.Generator().manual_seed(seed)
 
+    image = pipe(
+        prompt=prompt,
+        guidance_scale=0,
+        num_inference_steps=NUM_INFERENCE_STEPS,
+        generator=generator,
+    ).images[0]
 
-@spaces.GPU
-def predict(
-    input_image,
-    prompt,
-    negative_prompt,
-    seed,
-    guidance_scale=8.5,
-    controlnet_conditioning_scale=0.5,
-    strength=1.0,
-    controlnet_start=0.0,
-    controlnet_end=1.0,
-    guassian_sigma=2.0,
-    intensity_threshold=3,
-    progress=gr.Progress(track_tqdm=True),
-):
-    if input_image is None:
-        raise gr.Error("Please upload an image.")
-    padded_image = pad_image(input_image).resize((1024, 1024)).convert("RGB")
-    conditioning, pooled = compel([prompt, negative_prompt])
-    generator = torch.manual_seed(seed)
-    last_time = time.time()
-    anyline_image = anyline(
-        padded_image,
-        detect_resolution=1280,
-        guassian_sigma=max(0.01, guassian_sigma),
-        intensity_threshold=intensity_threshold,
-    )
+    return image
 
-    images = pipe(
-        image=padded_image,
-        control_image=anyline_image,
-        strength=strength,
-        prompt_embeds=conditioning[0:1],
-        pooled_prompt_embeds=pooled[0:1],
-        negative_prompt_embeds=conditioning[1:2],
-        negative_pooled_prompt_embeds=pooled[1:2],
-        width=1024,
-        height=1024,
-        controlnet_conditioning_scale=float(controlnet_conditioning_scale),
-        controlnet_start=float(controlnet_start),
-        controlnet_end=float(controlnet_end),
-        generator=generator,
-        num_inference_steps=30,
-        guidance_scale=guidance_scale,
-        eta=1.0,
-    )
-    print(f"Time taken: {time.time() - last_time}")
-    return (padded_image, images.images[0]), padded_image, anyline_image
 
+examples = [
+    "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.",
+    "A raccoon reading a book in a lush forest.",
+    "A small cactus with a happy face in the Sahara desert.",
+    "A super-realistic close-up of a snake eye",
+    "A cute cheetah looking amazed and surprised",
+    "Pirate ship sailing on a sea with the milky way galaxy in the sky and purple glow lights",
+    "a cute fluffy rabbit pilot walking on a military aircraft carrier, 8k, cinematic",
+    "A close up of an old elderly man with green eyes looking straight at the camera",
+    "A beautiful sunflower in rainy day",
+]
 
 css = """
-#intro{
-    # max-width: 32rem;
-    # text-align: center;
-    # margin: 0 auto;
+#col-container {
+    margin: 0 auto;
+    max-width: 512px;
 }
 """
 
+if torch.cuda.is_available():
+    power_device = "GPU"
+else:
+    power_device = "CPU"
+
 with gr.Blocks(css=css) as demo:
-    gr.Markdown(
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(
+            f"""
+        # ⚡ Flash Diffusion: FlashPixart ⚡
+        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 Conditional
+        Diffusion Model for Few Steps Image Generation](http://arxiv.org/abs/2406.02347) *by Clément Chadebec, Onur Tasar, Eyal Benaroche and Benjamin Aubin.*
+        [This model](https://huggingface.co/jasperai/flash-pixart) is a **66.5M** LoRA distilled version of [Pixart-α](https://huggingface.co/PixArt-alpha/PixArt-XL-2-1024-MS) model that is able to generate 1024x1024 images in **4 steps**. 
+        Currently running on {power_device}.
         """
-# MistoLine ControlNet demo
- 
-You can upload an initial image and prompt to generate an enhanced version.  
-SDXL Controlnet [TheMistoAI/MistoLine](https://huggingface.co/TheMistoAI/MistoLine)  
-[Anyline with Controlnet Aux ](https://github.com/huggingface/controlnet_aux)  
-For upscaling see [Enhance This Demo](https://huggingface.co/spaces/radames/Enhance-This-HiDiffusion-SDXL)  
-        """,
-        elem_id="intro",
-    )
-    with gr.Row():
-        with gr.Column(scale=1):
-            image_input = gr.Image(type="pil", label="Input Image")
-            prompt = gr.Textbox(
+        )
+        gr.Markdown(
+            "💡 *Hint:* To better appreciate the low latency of our method, run the demo locally !"
+        )
+
+        with gr.Row():
+            prompt = gr.Text(
                 label="Prompt",
-                info="The prompt is very important to get the desired results. Please try to describe the image as best as you can. Accepts Compel Syntax",
-            )
-            negative_prompt = gr.Textbox(
-                label="Negative Prompt",
-                value="blurry, ugly, duplicate, poorly drawn, deformed, mosaic",
+                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=2**64 - 1,
-                value=1415926535897932,
+                maximum=MAX_SEED,
                 step=1,
-                label="Seed",
-                randomize=True,
+                value=0,
             )
-            with gr.Accordion(label="Advanced", open=False):
-                guidance_scale = gr.Slider(
-                    minimum=0,
-                    maximum=50,
-                    value=8.5,
-                    step=0.001,
-                    label="Guidance Scale",
-                )
-                controlnet_conditioning_scale = gr.Slider(
-                    minimum=0,
-                    maximum=1,
-                    step=0.001,
-                    value=0.5,
-                    label="ControlNet Conditioning Scale",
-                )
-                strength = gr.Slider(
-                    minimum=0,
-                    maximum=1,
-                    step=0.001,
-                    value=1,
-                    label="Strength",
-                )
-                controlnet_start = gr.Slider(
-                    minimum=0,
-                    maximum=1,
-                    step=0.001,
-                    value=0.0,
-                    label="ControlNet Start",
-                )
-                controlnet_end = gr.Slider(
-                    minimum=0.0,
-                    maximum=1.0,
-                    step=0.001,
-                    value=1.0,
-                    label="ControlNet End",
-                )
-                guassian_sigma = gr.Slider(
-                    minimum=0.01,
-                    maximum=10.0,
-                    step=0.1,
-                    value=2.0,
-                    label="(Anyline) Guassian Sigma",
-                )
-                intensity_threshold = gr.Slider(
-                    minimum=0,
-                    maximum=255,
-                    step=1,
-                    value=3,
-                    label="(Anyline) Intensity Threshold",
-                )
 
-            btn = gr.Button()
-        with gr.Column(scale=2):
-            with gr.Group():
-                image_slider = ImageSlider(position=0.5)
-            with gr.Row():
-                padded_image = gr.Image(type="pil", label="Padded Image")
-                anyline_image = gr.Image(type="pil", label="Anyline Image")
-    inputs = [
-        image_input,
-        prompt,
-        negative_prompt,
-        seed,
-        guidance_scale,
-        controlnet_conditioning_scale,
-        strength,
-        controlnet_start,
-        controlnet_end,
-        guassian_sigma,
-        intensity_threshold,
-    ]
-    outputs = [image_slider, padded_image, anyline_image]
-    btn.click(lambda x: None, inputs=None, outputs=image_slider).then(
-        fn=predict, inputs=inputs, outputs=outputs
-    )
-    gr.Examples(
-        fn=predict,
-        inputs=inputs,
-        outputs=outputs,
-        examples=[
-            [
-                "./examples/city.png",
-                "hyperrealistic surreal cityscape scene at sunset, buildings",
-                "blurry, ugly, duplicate, poorly drawn, deformed, mosaic",
-                13113544138610326000,
-                8.5,
-                0.481,
-                1.0,
-                0.0,
-                0.9,
-                2,
-                3,
-            ],
-            [
-                "./examples/lara.jpeg",
-                "photography of lara croft 8k high definition award winning",
-                "blurry, ugly, duplicate, poorly drawn, deformed, mosaic",
-                5436236241,
-                8.5,
-                0.8,
-                1.0,
-                0.0,
-                0.9,
-                2,
-                3,
-            ],
-            [
-                "./examples/cybetruck.jpeg",
-                "photo of tesla cybertruck futuristic car 8k high definition on a sand dune in mars, future",
-                "blurry, ugly, duplicate, poorly drawn, deformed, mosaic",
-                383472451451,
-                8.5,
-                0.8,
-                0.8,
-                0.0,
-                0.9,
-                2,
-                3,
-            ],
-            [
-                "./examples/jesus.png",
-                "a photorealistic painting of Jesus Christ, 4k high definition",
-                "blurry, ugly, duplicate, poorly drawn, deformed, mosaic",
-                13317204146129588000,
-                8.5,
-                0.8,
-                0.8,
-                0.0,
-                0.9,
-                2,
-                3,
-            ],
-            [
-                "./examples/anna-sullivan-DioLM8ViiO8-unsplash.jpg",
-                "A crowded stadium with enthusiastic fans watching a daytime sporting event, the stands filled with colorful attire and the sun casting a warm glow",
-                "blurry, ugly, duplicate, poorly drawn, deformed, mosaic",
-                5623124123512,
-                8.5,
-                0.8,
-                0.8,
-                0.0,
-                0.9,
-                2,
-                3,
-            ],
-            [
-                "./examples/img_aef651cb-2919-499d-aa49-6d4e2e21a56e_1024.jpg",
-                "a large red flower on a black background 4k high definition",
-                "blurry, ugly, duplicate, poorly drawn, deformed, mosaic",
-                23123412341234,
-                8.5,
-                0.8,
-                0.8,
-                0.0,
-                0.9,
-                2,
-                3,
-            ],
-            [
-                "./examples/huggingface.jpg",
-                "photo realistic huggingface human emoji costume, round, yellow, (human skin)+++ (human texture)+++",
-                "blurry, ugly, duplicate, poorly drawn, deformed, mosaic, emoji cartoon,  drawing, pixelated",
-                12312353423,
-                15.206,
-                0.364,
-                0.8,
-                0.0,
-                0.9,
-                2,
-                3,
-            ],
-        ],
-        cache_examples="lazy",
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+
+        examples = gr.Examples(examples=examples, inputs=[prompt])
+
+        gr.Markdown("**Disclaimer:**")
+        gr.Markdown(
+            "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."
+        )
+    gr.on(
+        [run_button.click, seed.change, prompt.change, randomize_seed.change],
+        fn=infer,
+        inputs=[prompt, seed, randomize_seed],
+        outputs=[result],
+        show_progress="minimal",
+        show_api=False,
+        trigger_mode="always_last",
     )
 
-
-demo.queue(api_open=True)
-demo.launch(show_api=True)
+demo.queue().launch(show_api=False)