add SDXL-Lightning

app.py

@@ -1,23 +1,20 @@
 import os
 import shutil
 import tempfile
-import time
-from os import path
 
 import gradio as gr
 import numpy as np
 import rembg
 import spaces
 import torch
-from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler, StableDiffusionXLPipeline, LCMScheduler
+from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler, StableDiffusionXLPipeline, EulerDiscreteScheduler
 from einops import rearrange
 from huggingface_hub import hf_hub_download
 from omegaconf import OmegaConf
 from PIL import Image
 from pytorch_lightning import seed_everything
-from safetensors.torch import load_file
 from torchvision.transforms import v2
-from tqdm import tqdm
+from safetensors.torch import load_file
 
 from src.utils.camera_util import (FOV_to_intrinsics, get_circular_camera_poses,
                                    get_zero123plus_input_cameras)
@@ -25,7 +22,6 @@ from src.utils.infer_util import (remove_background, resize_foreground)
 from src.utils.mesh_util import save_glb, save_obj
 from src.utils.train_util import instantiate_from_config
 
-torch.backends.cuda.matmul.allow_tf32 = True
 
 def find_cuda():
     cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH')
@@ -130,18 +126,18 @@ def make3d(images):
 
 
 @spaces.GPU
-def process_image(num_images, prompt):
-    global pipe
-    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
-        return pipe(
-            prompt=[prompt]*num_images,
-            generator=torch.Generator().manual_seed(123),
-            num_inference_steps=1,
-            guidance_scale=0.,
-            height=int(512),
-            width=int(512),
-            timesteps=[800]
-        ).images
+def generate_image(prompt):
+    checkpoint = "sdxl_lightning_8step_unet.safetensors"
+    num_inference_steps = 8
+
+    pipe.scheduler = EulerDiscreteScheduler.from_config(
+        pipe.scheduler.config, timestep_spacing="trailing")
+    pipe.unet.load_state_dict(
+        load_file(hf_hub_download(repo, checkpoint), device="cuda"))
+
+    results = pipe(
+        prompt, num_inference_steps=num_inference_steps, guidance_scale=0)
+    return results.images[0]
 
 
 # Configuration
@@ -185,23 +181,24 @@ model.load_state_dict(state_dict, strict=True)
 
 model = model.to(device)
 
-# # Load text-to-image model
-# print('Loading text-to-image model ...')
+# Load StableDiffusionXL model
+base = "stabilityai/stable-diffusion-xl-base-1.0"
+repo = "ByteDance/SDXL-Lightning"
 
-# pipe = StableDiffusionXLPipeline.from_pretrained(
-#     "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16)
-# pipe.to(device="cuda", dtype=torch.bfloat16)
-
-# unet_state = load_file(hf_hub_download(
-#     "ByteDance/Hyper-SD", "Hyper-SDXL-1step-Unet.safetensors"), device="cuda")
-# pipe.unet.load_state_dict(unet_state)
-# pipe.scheduler = LCMScheduler.from_config(
-#     pipe.scheduler.config, timestep_spacing="trailing")
+pipe = StableDiffusionXLPipeline.from_pretrained(
+    base, torch_dtype=torch.float16, variant="fp16").to("cuda")
 
 print('Loading Finished!')
 
 # Gradio UI
 with gr.Blocks() as demo:
+    with gr.Group():
+        with gr.Row():
+            prompt = gr.Textbox(label='Enter your prompt (English)', scale=8)
+            submit_prompt = gr.Button(
+                scale=1, variant='primary', label='Generate Image')
+            img = gr.Image(label='SDXL-Lightning Generated Image')
+
     with gr.Row(variant="panel"):
         with gr.Column():
             with gr.Row():
@@ -228,7 +225,7 @@ with gr.Blocks() as demo:
                         label="Sample Steps", minimum=30, maximum=75, value=75, step=5)
 
             with gr.Row():
-                submit = gr.Button(
+                submit_mesh = gr.Button(
                     "Generate", elem_id="generate", variant="primary")
 
             with gr.Row(variant="panel"):
@@ -273,7 +270,9 @@ with gr.Blocks() as demo:
 
     mv_images = gr.State()
 
-    submit.click(fn=check_input_image, inputs=[input_image]).success(
+    submit_prompt.click(fn=generate_image, inputs=[prompt], outputs=img)
+
+    submit_mesh.click(fn=check_input_image, inputs=[input_image]).success(
         fn=preprocess,
         inputs=[input_image, do_remove_background],
         outputs=[processed_image],