init sdx

app.py

@@ -1,18 +1,21 @@
 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
+from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler, StableDiffusionXLPipeline, LCMScheduler
 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
 
@@ -22,6 +25,26 @@ 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
 
+cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
+os.environ["TRANSFORMERS_CACHE"] = cache_path
+os.environ["HF_HUB_CACHE"] = cache_path
+os.environ["HF_HOME"] = cache_path
+
+torch.backends.cuda.matmul.allow_tf32 = True
+
+
+class timer:
+    def __init__(self, method_name="timed process"):
+        self.method = method_name
+
+    def __enter__(self):
+        self.start = time.time()
+        print(f"{self.method} starts")
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        end = time.time()
+        print(f"{self.method} took {str(round(end - self.start, 2))}s")
+
 
 def find_cuda():
     cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH')
@@ -52,7 +75,7 @@ def get_render_cameras(batch_size=1, M=120, radius=2.5, elevation=10.0, is_flexi
 
 def check_input_image(input_image):
     if input_image is None:
-        raise gr.Error("No image uploaded!")
+        raise gr.Error("No image selected!")
 
 
 def preprocess(input_image, do_remove_background):
@@ -125,6 +148,21 @@ def make3d(images):
     return mesh_fpath, mesh_glb_fpath
 
 
+@spaces.GPU
+def process_image(num_images, prompt):
+    global pipe
+    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
+        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
+
+
 # Configuration
 cuda_path = find_cuda()
 config_path = 'configs/instant-mesh-large.yaml'
@@ -166,6 +204,21 @@ model.load_state_dict(state_dict, strict=True)
 
 model = model.to(device)
 
+# Load text-to-image model
+print('Loading text-to-image model ...')
+if not path.exists(cache_path):
+    os.makedirs(cache_path, exist_ok=True)
+
+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")
+
 print('Loading Finished!')
 
 # Gradio UI