Update app.py

app.py

@@ -1,15 +1,12 @@
 #!/usr/bin/env python
 
-from __future__ import annotations
-
 import os
 import random
-
 import gradio as gr
 import numpy as np
 import PIL.Image
 import torch
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, UNet2DConditionModel, LCMScheduler
 
 DESCRIPTION = 'This space is an API service meant to be used by VideoChain and VideoQuest.\nWant to use this space for yourself? Please use the original code: [https://huggingface.co/spaces/hysts/SD-XL](https://huggingface.co/spaces/hysts/SD-XL)'
 if not torch.cuda.is_available():
@@ -17,37 +14,27 @@ if not torch.cuda.is_available():
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = int(os.getenv('MAX_IMAGE_SIZE', '1024'))
-USE_TORCH_COMPILE = os.getenv('USE_TORCH_COMPILE') == '1'
-ENABLE_CPU_OFFLOAD = os.getenv('ENABLE_CPU_OFFLOAD') == '1'
 SECRET_TOKEN = os.getenv('SECRET_TOKEN', 'default_secret')
 
 device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
 if torch.cuda.is_available():
-    pipe = DiffusionPipeline.from_pretrained(
-        'stabilityai/stable-diffusion-xl-base-1.0',
-        torch_dtype=torch.float16,
-        use_safetensors=True,
-        variant='fp16')
-    refiner = DiffusionPipeline.from_pretrained(
-        'stabilityai/stable-diffusion-xl-refiner-1.0',
+    unet = UNet2DConditionModel.from_pretrained(
+        "latent-consistency/lcm-ssd-1b",
         torch_dtype=torch.float16,
-        use_safetensors=True,
-        variant='fp16')
+        variant="fp16"
+    )
 
-    if ENABLE_CPU_OFFLOAD:
-        pipe.enable_model_cpu_offload()
-        refiner.enable_model_cpu_offload()
-    else:
-        pipe.to(device)
-        refiner.to(device)
+    pipe = DiffusionPipeline.from_pretrained(
+        "segmind/SSD-1B",
+        unet=unet,
+        torch_dtype=torch.float16,
+        variant="fp16"
+    )
 
-    if USE_TORCH_COMPILE:
-        pipe.unet = torch.compile(pipe.unet,
-                                  mode='reduce-overhead',
-                                  fullgraph=True)
+    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+    pipe.to(device)
 else:
     pipe = None
-    refiner = None
     
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
@@ -65,11 +52,8 @@ def generate(prompt: str,
              seed: int = 0,
              width: int = 1024,
              height: int = 1024,
-             guidance_scale_base: float = 5.0,
-             guidance_scale_refiner: float = 5.0,
-             num_inference_steps_base: int = 50,
-             num_inference_steps_refiner: int = 50,
-             apply_refiner: bool = False,
+             guidance_scale: float = 1.0,
+             num_inference_steps: int = 4,
              secret_token: str = '') -> PIL.Image.Image:
     if secret_token != SECRET_TOKEN:
         raise gr.Error(
@@ -84,37 +68,16 @@ def generate(prompt: str,
     if not use_negative_prompt_2:
         negative_prompt_2 = None  # type: ignore
 
-    if not apply_refiner:
-        return pipe(prompt=prompt,
-                    negative_prompt=negative_prompt,
-                    prompt_2=prompt_2,
-                    negative_prompt_2=negative_prompt_2,
-                    width=width,
-                    height=height,
-                    guidance_scale=guidance_scale_base,
-                    num_inference_steps=num_inference_steps_base,
-                    generator=generator,
-                    output_type='pil').images[0]
-    else:
-        latents = pipe(prompt=prompt,
-                       negative_prompt=negative_prompt,
-                       prompt_2=prompt_2,
-                       negative_prompt_2=negative_prompt_2,
-                       width=width,
-                       height=height,
-                       guidance_scale=guidance_scale_base,
-                       num_inference_steps=num_inference_steps_base,
-                       generator=generator,
-                       output_type='latent').images
-        image = refiner(prompt=prompt,
-                        negative_prompt=negative_prompt,
-                        prompt_2=prompt_2,
-                        negative_prompt_2=negative_prompt_2,
-                        guidance_scale=guidance_scale_refiner,
-                        num_inference_steps=num_inference_steps_refiner,
-                        image=latents,
-                        generator=generator).images[0]
-        return image
+    return pipe(prompt=prompt,
+                negative_prompt=negative_prompt,
+                prompt_2=prompt_2,
+                negative_prompt_2=negative_prompt_2,
+                width=width,
+                height=height,
+                guidance_scale=guidance_scale,
+                num_inference_steps=num_inference_steps,
+                generator=generator,
+                output_type='pil').images[0]
 
 with gr.Blocks(css='style.css') as demo:
     gr.Markdown(DESCRIPTION)
@@ -181,61 +144,33 @@ with gr.Blocks(css='style.css') as demo:
                     step=32,
                     value=1024,
                 )
-            apply_refiner = gr.Checkbox(label='Apply refiner', value=False)
-            with gr.Row():
-                guidance_scale_base = gr.Slider(
-                    label='Guidance scale for base',
-                    minimum=1,
-                    maximum=20,
-                    step=0.1,
-                    value=5.0)
-                num_inference_steps_base = gr.Slider(
-                    label='Number of inference steps for base',
-                    minimum=10,
-                    maximum=100,
-                    step=1,
-                    value=50)
-            with gr.Row(visible=False) as refiner_params:
-                guidance_scale_refiner = gr.Slider(
-                    label='Guidance scale for refiner',
-                    minimum=1,
-                    maximum=20,
-                    step=0.1,
-                    value=5.0)
-                num_inference_steps_refiner = gr.Slider(
-                    label='Number of inference steps for refiner',
-                    minimum=10,
-                    maximum=100,
-                    step=1,
-                    value=50)
+            guidance_scale = gr.Slider(
+                label='Guidance scale',
+                minimum=1,
+                maximum=20,
+                step=0.1,
+                value=1.0)
+            num_inference_steps = gr.Slider(
+                label='Number of inference steps',
+                minimum=2,
+                maximum=8,
+                step=1,
+                value=4)
 
     use_negative_prompt.change(
         fn=lambda x: gr.update(visible=x),
         inputs=use_negative_prompt,
-        outputs=negative_prompt,
-        queue=False,
-        api_name=False,
+        outputs=negative_prompt
     )
     use_prompt_2.change(
         fn=lambda x: gr.update(visible=x),
         inputs=use_prompt_2,
-        outputs=prompt_2,
-        queue=False,
-        api_name=False,
+        outputs=prompt_2
     )
     use_negative_prompt_2.change(
         fn=lambda x: gr.update(visible=x),
         inputs=use_negative_prompt_2,
-        outputs=negative_prompt_2,
-        queue=False,
-        api_name=False,
-    )
-    apply_refiner.change(
-        fn=lambda x: gr.update(visible=x),
-        inputs=apply_refiner,
-        outputs=refiner_params,
-        queue=False,
-        api_name=False,
+        outputs=negative_prompt_2
     )
 
     inputs = [
@@ -249,19 +184,14 @@ with gr.Blocks(css='style.css') as demo:
         seed,
         width,
         height,
-        guidance_scale_base,
-        guidance_scale_refiner,
-        num_inference_steps_base,
-        num_inference_steps_refiner,
-        apply_refiner,
+        guidance_scale,
+        num_inference_steps,
         secret_token,
     ]
     prompt.submit(
         fn=randomize_seed_fn,
         inputs=[seed, randomize_seed],
-        outputs=seed,
-        queue=False,
-        api_name=False,
+        outputs=seed
     ).then(
         fn=generate,
         inputs=inputs,
@@ -271,25 +201,19 @@ with gr.Blocks(css='style.css') as demo:
     negative_prompt.submit(
         fn=randomize_seed_fn,
         inputs=[seed, randomize_seed],
-        outputs=seed,
-        queue=False,
-        api_name=False,
+        outputs=seed
     ).then(
         fn=generate,
         inputs=inputs,
-        outputs=result,
-        api_name=False,
+        outputs=result
     )
     run_button.click(
         fn=randomize_seed_fn,
         inputs=[seed, randomize_seed],
-        outputs=seed,
-        queue=False,
-        api_name=False,
+        outputs=seed
     ).then(
         fn=generate,
         inputs=inputs,
-        outputs=result,
-        api_name=False,
+        outputs=result
     )
-demo.queue(max_size=6).launch()
+demo.queue(max_size=6).launch()