remove sdxl

app.py

@@ -1,21 +1,18 @@
 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
 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
 
@@ -25,26 +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
 
-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')
@@ -75,7 +52,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 selected!")
+        raise gr.Error("No image uploaded!")
 
 
 def preprocess(input_image, do_remove_background):
@@ -148,21 +125,6 @@ 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'
@@ -204,21 +166,6 @@ 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
@@ -226,23 +173,19 @@ with gr.Blocks() as demo:
     with gr.Row(variant="panel"):
         with gr.Column():
             with gr.Row():
-                num_images = gr.Slider(
-                    label="Number of Images", minimum=1, maximum=8, step=1, value=4, interactive=True)
-                prompt = gr.Text(
-                    label="Prompt", value="a photo of a cat", interactive=True)
-                generate_2d_btn = gr.Button(value="Generate 2D Images")
-
-            with gr.Row():
-                generated_images = gr.Gallery(height=512)
-
-            with gr.Row():
-                selected_image = gr.Image(
-                    label="Selected Image",
+                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",
                     image_mode="RGBA",
                     type="pil",
                     interactive=False
                 )
-
             with gr.Row():
                 with gr.Group():
                     do_remove_background = gr.Checkbox(
@@ -253,8 +196,18 @@ with gr.Blocks() as demo:
                         label="Sample Steps", minimum=30, maximum=75, value=75, step=5)
 
             with gr.Row():
-                generate_3d_btn = gr.Button(
-                    "Generate 3D Model", elem_id="generate", variant="primary")
+                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
+                )
 
         with gr.Column():
             with gr.Row():
@@ -288,24 +241,13 @@ with gr.Blocks() as demo:
 
     mv_images = gr.State()
 
-    generate_2d_btn.click(
-        fn=process_image,
-        inputs=[num_images, prompt],
-        outputs=[generated_images]
-    )
-
-    def select_image(evt: gr.SelectData):
-        return evt.value['image']['url']
-
-    generated_images.select(select_image, None, selected_image)
-
-    generate_3d_btn.click(fn=check_input_image, inputs=[selected_image]).success(
+    submit.click(fn=check_input_image, inputs=[input_image]).success(
         fn=preprocess,
-        inputs=[selected_image, do_remove_background],
-        outputs=[selected_image],
+        inputs=[input_image, do_remove_background],
+        outputs=[processed_image],
     ).success(
         fn=generate_mvs,
-        inputs=[selected_image, sample_steps, sample_seed],
+        inputs=[processed_image, sample_steps, sample_seed],
         outputs=[mv_images, mv_show_images]
     ).success(
         fn=make3d,
@@ -313,4 +255,4 @@ with gr.Blocks() as demo:
         outputs=[output_model_obj, output_model_glb]
     )
 
-demo.launch()
+demo.launch()