prompt PromptWeighting

README.md

@@ -4,7 +4,7 @@ emoji: ⚡️⚡️⚡️⚡️
 colorFrom: yellow
 colorTo: purple
 sdk: gradio
-sdk_version: 4.19.2
+sdk_version: 4.25.0
 app_file: app.py
 pinned: false
 disable_embedding: true

app.py

@@ -1,3 +1,4 @@
+import spaces
 from diffusers import (
     StableDiffusionXLPipeline,
     EulerDiscreteScheduler,
@@ -7,6 +8,8 @@ from diffusers import (
 import torch
 import os
 from huggingface_hub import hf_hub_download
+from compel import Compel, ReturnedEmbeddingsType
+from gradio_promptweighting import PromptWeighting
 
 
 from PIL import Image
@@ -24,13 +27,6 @@ REPO = "ByteDance/SDXL-Lightning"
 CHECKPOINT = "sdxl_lightning_2step_unet.safetensors"
 taesd_model = "madebyollin/taesdxl"
 
-# {
-#     "1-Step": ["sdxl_lightning_1step_unet_x0.safetensors", 1],
-#     "2-Step": ["sdxl_lightning_2step_unet.safetensors", 2],
-#     "4-Step": ["sdxl_lightning_4step_unet.safetensors", 4],
-#     "8-Step": ["sdxl_lightning_8step_unet.safetensors", 8],
-# }
-
 
 SFAST_COMPILE = os.environ.get("SFAST_COMPILE", "0") == "1"
 SAFETY_CHECKER = os.environ.get("SAFETY_CHECKER", "0") == "1"
@@ -55,6 +51,15 @@ unet.load_state_dict(load_file(hf_hub_download(REPO, CHECKPOINT), device="cuda")
 pipe = StableDiffusionXLPipeline.from_pretrained(
     BASE, unet=unet, torch_dtype=torch.float16, variant="fp16", safety_checker=False
 ).to("cuda")
+unet = unet.to(dtype=torch.float16)
+
+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],
+)
+
 
 if USE_TAESD:
     pipe.vae = AutoencoderTiny.from_pretrained(
@@ -115,14 +120,24 @@ if SFAST_COMPILE:
     pipe = compile(pipe, config)
 
 
-def predict(prompt, seed=1231231):
+@spaces.GPU
+def predict(prompt, prompt_w, guidance_scale, seed=1231231):
     generator = torch.manual_seed(seed)
     last_time = time.time()
+    prompt_w = " ".join(
+        [f"({p['prompt']}){p['scale']}" for p in prompt_w if p["prompt"]]
+    )
+
+    conditioning, pooled = compel([prompt + " " + prompt_w, ""])
+
     results = pipe(
-        prompt=prompt,
+        prompt_embeds=conditioning[0:1],
+        pooled_prompt_embeds=pooled[0:1],
+        negative_prompt_embeds=conditioning[1:2],
+        negative_pooled_prompt_embeds=pooled[1:2],
         generator=generator,
         num_inference_steps=2,
-        guidance_scale=0.0,
+        guidance_scale=guidance_scale,
         # width=768,
         # height=768,
         output_type="pil",
@@ -142,12 +157,15 @@ def predict(prompt, seed=1231231):
 css = """
 #container{
     margin: 0 auto;
-    max-width: 40rem;
+    max-width: 80rem;
 }
 #intro{
     max-width: 100%;
     margin: 0 auto;
 }
+.generating {
+    display: none
+}
 """
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="container"):
@@ -159,51 +177,61 @@ with gr.Blocks(css=css) as demo:
             elem_id="intro",
         )
         with gr.Row():
-            with gr.Row():
-                prompt = gr.Textbox(
-                    placeholder="Insert your prompt here:", scale=5, container=False
-                )
-                generate_bt = gr.Button("Generate", scale=1)
-
-        image = gr.Image(type="filepath")
-        with gr.Accordion("Advanced options", open=False):
-            seed = gr.Slider(
-                randomize=True, minimum=0, maximum=12013012031030, label="Seed", step=1
-            )
-        with gr.Accordion("Run with diffusers"):
-            gr.Markdown(
-                """## Running SDXL-Lightning with `diffusers`
-```py
-import torch
-from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
-from huggingface_hub import hf_hub_download
-from safetensors.torch import load_file
-
-base = "stabilityai/stable-diffusion-xl-base-1.0"
-repo = "ByteDance/SDXL-Lightning"
-ckpt = "sdxl_lightning_2step_unet.safetensors" # Use the correct ckpt for your step setting!
-
-# Load model.
-unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
-unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cuda"))
-pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16, variant="fp16").to("cuda")
-
-# Ensure sampler uses "trailing" timesteps.
-pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
-
-# Ensure using the same inference steps as the loaded model and CFG set to 0.
-pipe("A girl smiling", num_inference_steps=2, guidance_scale=0).images[0].save("output.png")
-```
-            """
-            )
-
-        inputs = [prompt, seed]
+            with gr.Column():
+                with gr.Group():
+                    prompt = gr.Textbox(
+                        placeholder="Insert your prompt here:",
+                        max_lines=1,
+                        label="Prompt",
+                    )
+                    prompt_w = PromptWeighting(
+                        min=0,
+                        max=3,
+                        step=0.005,
+                        show_label=False,
+                    )
+
+                with gr.Accordion("Advanced options", open=True):
+                    seed = gr.Slider(
+                        minimum=0,
+                        maximum=12013012031030,
+                        label="Seed",
+                        step=1,
+                    )
+                    guidance_scale = gr.Slider(
+                        minimum=0.0,
+                        maximum=20.0,
+                        label="Guidance scale",
+                        value=0.0,
+                        step=0.1,
+                    )
+                generate_bt = gr.Button("Generate")
+            with gr.Column():
+                image = gr.Image(type="filepath")
+
+        inputs = [
+            prompt,
+            prompt_w,
+            guidance_scale,
+            seed,
+        ]
         outputs = [image]
-        generate_bt.click(
-            fn=predict, inputs=inputs, outputs=outputs, show_progress=False
+
+        gr.on(
+            triggers=[
+                prompt.input,
+                prompt_w.input,
+                generate_bt.click,
+                guidance_scale.input,
+                seed.input,
+            ],
+            fn=predict,
+            inputs=inputs,
+            outputs=outputs,
+            show_progress="hidden",
+            show_api=False,
+            trigger_mode="always_last",
         )
-        prompt.input(fn=predict, inputs=inputs, outputs=outputs, show_progress=False)
-        seed.change(fn=predict, inputs=inputs, outputs=outputs, show_progress=False)
 
-demo.queue()
+demo.queue(api_open=False)
 demo.launch()

requirements.txt

@@ -1,6 +1,6 @@
-diffusers==0.26.3
+diffusers==0.27.2
 transformers
-gradio==4.19.2
+gradio==4.25.0
 torch==2.1.0
 fastapi==0.104.0
 uvicorn==0.23.2
@@ -13,4 +13,7 @@ xformers
 hf_transfer
 huggingface_hub
 safetensors
-stable_fast @ https://github.com/chengzeyi/stable-fast/releases/download/v1.0.2/stable_fast-1.0.2+torch211cu121-cp310-cp310-manylinux2014_x86_64.whl; sys_platform != 'darwin' or platform_machine != 'arm64'
+compel
+stable_fast @ https://github.com/chengzeyi/stable-fast/releases/download/v1.0.2/stable_fast-1.0.2+torch211cu121-cp310-cp310-manylinux2014_x86_64.whl ; sys_platform != 'darwin' or platform_machine != 'arm64'
+spaces
+gradio_promptweighting @ ./gradio_promptweighting-0.0.1-py3-none-any.whl