add randomize seed and sfast

app.py

@@ -12,10 +12,13 @@ from PIL import Image
 import numpy as np
 import gradio as gr
 import psutil
+from sfast.compilers.stable_diffusion_pipeline_compiler import (
+    compile,
+    CompilationConfig,
+)
 
 
 SAFETY_CHECKER = os.environ.get("SAFETY_CHECKER", None)
-TORCH_COMPILE = os.environ.get("TORCH_COMPILE", None)
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 # check if MPS is available OSX only M1/M2/M3 chips
 mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
@@ -27,7 +30,6 @@ torch_device = device
 torch_dtype = torch.float16
 
 print(f"SAFETY_CHECKER: {SAFETY_CHECKER}")
-print(f"TORCH_COMPILE: {TORCH_COMPILE}")
 print(f"device: {device}")
 
 if mps_available:
@@ -43,24 +45,21 @@ else:
     pipe = DiffusionPipeline.from_pretrained(model_id, safety_checker=None)
 
 pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
-pipe.to(device=torch_device, dtype=torch_dtype).to(device)
-pipe.unet.to(memory_format=torch.channels_last)
-
-# check if computer has less than 64GB of RAM using sys or os
-if psutil.virtual_memory().total < 64 * 1024**3:
-    pipe.enable_attention_slicing()
-
-if TORCH_COMPILE:
-    pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
-    pipe.vae = torch.compile(pipe.vae, mode="reduce-overhead", fullgraph=True)
-
-    pipe(prompt="warmup", num_inference_steps=1, guidance_scale=8.0)
-
-# Load LCM LoRA
 pipe.load_lora_weights(
     "latent-consistency/lcm-lora-sdxl",
     use_auth_token=HF_TOKEN,
 )
+if device.type != "mps":
+    pipe.unet.to(memory_format=torch.channels_last)
+pipe.to(device=torch_device, dtype=torch_dtype).to(device)
+
+# Load LCM LoRA
+
+config = CompilationConfig.Default()
+config.enable_xformers = True
+config.enable_triton = True
+config.enable_cuda_graph = True
+pipe = compile(pipe, config=config)
 
 compel_proc = Compel(
     tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
@@ -71,8 +70,15 @@ compel_proc = Compel(
 
 
 def predict(
-    prompt, guidance, steps, seed=1231231, progress=gr.Progress(track_tqdm=True)
+    prompt,
+    guidance,
+    steps,
+    seed=1231231,
+    randomize_bt=False,
+    progress=gr.Progress(track_tqdm=True),
 ):
+    if randomize_bt:
+        seed = np.random.randint(0, 2**32 - 1)
     generator = torch.manual_seed(seed)
     prompt_embeds, pooled_prompt_embeds = compel_proc(prompt)
 
@@ -94,7 +100,7 @@ def predict(
     )
     if nsfw_content_detected:
         raise gr.Error("NSFW content detected.")
-    return results.images[0]
+    return results.images[0], seed
 
 
 css = """
@@ -122,18 +128,28 @@ with gr.Blocks(css=css) as demo:
                     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):
             guidance = gr.Slider(
                 label="Guidance", minimum=0.0, maximum=5, value=0.3, step=0.001
             )
             steps = gr.Slider(label="Steps", value=4, minimum=2, maximum=10, step=1)
-            seed = gr.Slider(
-                randomize=True, minimum=0, maximum=12013012031030, label="Seed", step=1
-            )
+            with gr.Row():
+                seed = gr.Slider(
+                    randomize=True,
+                    minimum=0,
+                    maximum=12013012031030,
+                    label="Seed",
+                    step=1,
+                    scale=5,
+                )
+                with gr.Group():
+                    randomize_bt = gr.Checkbox(label="Randomize", value=False)
+                    random_seed = gr.Textbox(show_label=False)
         with gr.Accordion("Run with diffusers"):
-            gr.Markdown('''## Running LCM-LoRAs it with `diffusers`
+            gr.Markdown(
+                """## Running LCM-LoRAs it with `diffusers`
             ```bash
             pip install diffusers==0.23.0
             ```
@@ -151,10 +167,11 @@ with gr.Blocks(css=css) as demo:
             )
             results.images[0]
             ```
-            ''')
-            
-        inputs = [prompt, guidance, steps, seed]
-        generate_bt.click(fn=predict, inputs=inputs, outputs=image)
+            """
+            )
+
+        inputs = [prompt, guidance, steps, seed, randomize_bt]
+        generate_bt.click(fn=predict, inputs=inputs, outputs=[image, random_seed])
 
 demo.queue()
 demo.launch()

requirements.txt

@@ -11,3 +11,6 @@ accelerate==0.24.0
 compel==2.0.2
 controlnet-aux==0.0.7
 peft==0.6.0
+stable_fast @ https://github.com/chengzeyi/stable-fast/releases/download/v0.0.15.post1/stable_fast-0.0.15.post1+torch211cu121-cp310-cp310-manylinux2014_x86_64.whl
+xformers
+triton