add sdxl

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, height, width, prompt, seed):
+    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(int(seed)),
+            num_inference_steps=1,
+            guidance_scale=0.,
+            height=int(height),
+            width=int(width),
+            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
@@ -173,19 +226,28 @@ with gr.Blocks() as demo:
     with gr.Row(variant="panel"):
         with gr.Column():
             with gr.Row():
-                input_image = gr.Image(
-                    label="Input Image",
-                    image_mode="RGBA",
-                    sources="upload",
-                    type="pil",
-                    elem_id="content_image",
-                )
-                processed_image = gr.Image(
-                    label="Processed Image",
+                num_images = gr.Slider(
+                    label="Number of Images", minimum=1, maximum=8, step=1, value=4, interactive=True)
+                height = gr.Number(label="Image Height",
+                                   value=1024, interactive=True)
+                width = gr.Number(label="Image Width",
+                                  value=1024, interactive=True)
+                prompt = gr.Text(
+                    label="Prompt", value="a photo of a cat", interactive=True)
+                seed = gr.Number(label="Seed", value=3413, interactive=True)
+                generate_2d_btn = gr.Button(value="Generate 2D Images")
+
+            with gr.Row():
+                generated_images = gr.Gallery(height=1024)
+
+            with gr.Row():
+                selected_image = gr.Image(
+                    label="Selected Image",
                     image_mode="RGBA",
                     type="pil",
                     interactive=False
                 )
+
             with gr.Row():
                 with gr.Group():
                     do_remove_background = gr.Checkbox(
@@ -196,18 +258,8 @@ with gr.Blocks() as demo:
                         label="Sample Steps", minimum=30, maximum=75, value=75, step=5)
 
             with gr.Row():
-                submit = gr.Button(
-                    "Generate", elem_id="generate", variant="primary")
-
-            with gr.Row(variant="panel"):
-                gr.Examples(
-                    examples=[os.path.join("examples", img_name)
-                              for img_name in sorted(os.listdir("examples"))],
-                    inputs=[input_image],
-                    label="Examples",
-                    cache_examples=False,
-                    examples_per_page=16
-                )
+                generate_3d_btn = gr.Button(
+                    "Generate 3D Model", elem_id="generate", variant="primary")
 
         with gr.Column():
             with gr.Row():
@@ -241,13 +293,25 @@ with gr.Blocks() as demo:
 
     mv_images = gr.State()
 
-    submit.click(fn=check_input_image, inputs=[input_image]).success(
+    generate_2d_btn.click(
+        fn=process_image,
+        inputs=[num_images, height, width, prompt, seed],
+        outputs=[generated_images]
+    )
+
+    generated_images.select(
+        fn=lambda x: x,
+        inputs=[generated_images],
+        outputs=[selected_image]
+    )
+
+    generate_3d_btn.click(fn=check_input_image, inputs=[selected_image]).success(
         fn=preprocess,
-        inputs=[input_image, do_remove_background],
-        outputs=[processed_image],
+        inputs=[selected_image, do_remove_background],
+        outputs=[selected_image],
     ).success(
         fn=generate_mvs,
-        inputs=[processed_image, sample_steps, sample_seed],
+        inputs=[selected_image, sample_steps, sample_seed],
         outputs=[mv_images, mv_show_images]
     ).success(
         fn=make3d,

requirements.txt

@@ -12,12 +12,17 @@ tensorboard
 PyMCubes
 trimesh
 rembg
-transformers==4.34.1
-diffusers==0.19.3
+transformers==4.38.2
+diffusers==0.25.0
 bitsandbytes
 imageio[ffmpeg]
 xatlas
 plyfile
 xformers==0.0.22.post7
 git+https://github.com/NVlabs/nvdiffrast/
-huggingface-hub
+huggingface-hub
+
+httpx==0.23.0
+flask
+pillow
+safetensors