server compatibility

app.py

@@ -12,7 +12,7 @@ from settings import ALLOW_CHANGING_BASE_MODEL, DEFAULT_MODEL_ID, SHOW_DUPLICATE
 DESCRIPTION = "# Material Authoring Demo v0.1. Under Construction"
 
 if not torch.cuda.is_available():
-    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
+    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p> Check if the 'CUDA_VISIBLE_DEVICES' are set incorrectly in settings.py"
 
 # model = Model(base_model_id=DEFAULT_MODEL_ID, task_name="Canny")
 model = Model(base_model_id=DEFAULT_MODEL_ID, task_name="texnet")

app_texnet.py

@@ -174,52 +174,52 @@ def create_demo(process):
                     "examples/bunny/frame_0001.png", # /dgxusers/Users/jyang/project/ObjectReal/data/control/preprocess/bunny/uv_normal/fused.png
                     "examples/bunny/uv_normal.png", # /dgxusers/Users/jyang/project/ObjectReal/data/control/preprocess/bunny/uv_normal/fused.png
                     "feather", 
-                    # a_prompt.value, 
-                    # n_prompt.value,
-                    # num_samples.value,
-                    # image_resolution.value,
-                    # num_steps.value,
-                    # guidance_scale.value,
-                    # seed.value,
+                    a_prompt.value, 
+                    n_prompt.value,
+                    num_samples.value,
+                    image_resolution.value,
+                    num_steps.value,
+                    guidance_scale.value,
+                    seed.value,
                 ],
                 [
                     "monkey",
                     "examples/monkey/frame_0001.png", # /dgxusers/Users/jyang/project/ObjectReal/data/control/preprocess/monkey/uv_normal/fused.png
                     "examples/monkey/uv_normal.png", # /dgxusers/Users/jyang/project/ObjectReal/data/control/preprocess/monkey/uv_normal/fused.png
                     "wood", 
-                    # a_prompt.value, 
-                    # n_prompt.value,
-                    # num_samples.value,
-                    # image_resolution.value,
-                    # num_steps.value,
-                    # guidance_scale.value,
-                    # seed.value,
+                    a_prompt.value, 
+                    n_prompt.value,
+                    num_samples.value,
+                    image_resolution.value,
+                    num_steps.value,
+                    guidance_scale.value,
+                    seed.value,
                 ],
                 [
                     "tshirt",
                     "examples/tshirt/frame_0001.png", # /dgxusers/Users/jyang/project/ObjectReal/data/control/preprocess/monkey/uv_normal/fused.png
                     "examples/tshirt/uv_normal.png", # /dgxusers/Users/jyang/project/ObjectReal/data/control/preprocess/monkey/uv_normal/fused.png
                     "wood", 
-                    # a_prompt.value, 
-                    # n_prompt.value,
-                    # num_samples.value,
-                    # image_resolution.value,
-                    # num_steps.value,
-                    # guidance_scale.value,
-                    # seed.value,
+                    a_prompt.value, 
+                    n_prompt.value,
+                    num_samples.value,
+                    image_resolution.value,
+                    num_steps.value,
+                    guidance_scale.value,
+                    seed.value,
                 ],
                 [
                     "highheel",
                     "examples/highheel/frame_0001.png", # /dgxusers/Users/jyang/project/ObjectReal/data/control/preprocess/monkey/uv_normal/fused.png
                     "examples/highheel/uv_normal.png", # /dgxusers/Users/jyang/project/ObjectReal/data/control/preprocess/monkey/uv_normal/fused.png
                     "wood", 
-                    # a_prompt.value, 
-                    # n_prompt.value,
-                    # num_samples.value,
-                    # image_resolution.value,
-                    # num_steps.value,
-                    # guidance_scale.value,
-                    # seed.value,
+                    a_prompt.value, 
+                    n_prompt.value,
+                    num_samples.value,
+                    image_resolution.value,
+                    num_steps.value,
+                    guidance_scale.value,
+                    seed.value,
                 ],
             ],
         )

model.py

@@ -66,6 +66,24 @@ class Model:
         ).to(self.device)
         self.preprocessor = Preprocessor()
 
+        # set up pipe_rgb2x
+        self.pipe_rgb2x = StableDiffusionAOVMatEstPipeline.from_pretrained(
+            "zheng95z/rgb-to-x",
+            torch_dtype=torch.float16,
+        ).to(self.device)
+        self.pipe_rgb2x.scheduler = DDIMScheduler.from_config(
+            self.pipe_rgb2x.scheduler.config, rescale_betas_zero_snr=True, timestep_spacing="trailing"
+        )
+        self.pipe_rgb2x.set_progress_bar_config(disable=True)
+
+        # setup blender
+        self.blender_path = '/tmp/blender-3.2.2-linux-x64/blender'
+        if not os.path.exists(self.blender_path):
+            print("Downloading Blender...")
+            subprocess.run(["wget", "https://download.blender.org/release/Blender3.2/blender-3.2.2-linux-x64.tar.xz", "-O", "/tmp/blender-3.2.2-linux-x64.tar.xz"], check=True)
+            subprocess.run(["tar", "-xf", "/tmp/blender-3.2.2-linux-x64.tar.xz", "-C", "/tmp"], check=True)
+            print("Blender downloaded and extracted.")
+
     def load_pipe(self, base_model_id: str, task_name: str) -> DiffusionPipeline:
         if (
             base_model_id == self.base_model_id
@@ -293,18 +311,7 @@ class Model:
             return photo, return_list, prompts
         
         # Load rgb2x pipeline
-        pipe = StableDiffusionAOVMatEstPipeline.from_pretrained(
-            "zheng95z/rgb-to-x",
-            torch_dtype=torch.float16,
-            # cache_dir=os.path.join(current_directory, "model_cache"),
-        ).to(self.pipe.device)
-        pipe.scheduler = DDIMScheduler.from_config(
-            pipe.scheduler.config, rescale_betas_zero_snr=True, timestep_spacing="trailing"
-        )
-        pipe.set_progress_bar_config(disable=True)
-        pipe.to(self.pipe.device)
-
-        _, preds, prompts = rgb2x(pipe, torchvision.transforms.PILToTensor()(tex_fine[0]).to(self.pipe.device), inference_step=num_steps, num_samples=num_images)
+        _, preds, prompts = rgb2x(self.pipe_rgb2x, torchvision.transforms.PILToTensor()(tex_fine[0]).to(self.pipe.device), inference_step=num_steps, num_samples=num_images)
 
         base_color_path = image_to_temp_path(tex_fine[0].rotate(90), "base_color")
         normal_map_path = image_to_temp_path(preds[0].rotate(90), "normal_map")
@@ -313,7 +320,6 @@ class Model:
         output_blend_path = os.path.join(os.getcwd(), "output", f"{obj_name}.blend")  # replace with desired output path
         os.makedirs(os.path.dirname(output_blend_path), exist_ok=True)
 
-        # 
         def run_blend_generation(
             blender_path,
             generate_script_path,
@@ -331,15 +337,8 @@ class Model:
             subprocess.run(cmd, check=True)
 
         # check if the blender_path exists, if not download
-        blender_path = '/tmp/blender-3.2.2-linux-x64/blender'
-        if not os.path.exists(blender_path):
-            print("Downloading Blender...")
-            subprocess.run(["wget", "https://download.blender.org/release/Blender3.2/blender-3.2.2-linux-x64.tar.xz", "-O", "/tmp/blender-3.2.2-linux-x64.tar.xz"], check=True)
-            subprocess.run(["tar", "-xf", "/tmp/blender-3.2.2-linux-x64.tar.xz", "-C", "/tmp"], check=True)
-            print("Blender downloaded and extracted.")
-
         run_blend_generation(
-            blender_path=blender_path,
+            blender_path=self.blender_path,
             generate_script_path="rgb2x/generate_blend.py",
             obj_path=f"examples/{obj_name}/mesh.obj",  # replace with actual mesh path
             base_color_path=base_color_path,
@@ -351,8 +350,6 @@ class Model:
 
         # gallary
         return [*tex_fine], [preds[1]], [preds[2]], [preds[3]], [output_blend_path]
-        # image
-        # return tex_fine[0], preds[1], preds[2], preds[3], output_blend_path
 
     # @spaces.GPU #[uncomment to use ZeroGPU]
     @torch.inference_mode()

settings.py

@@ -18,5 +18,5 @@ MAX_SEED = np.iinfo(np.int32).max
 
 # setup CUDA
 # disable the following when deployting to hugging face
-if os.getenv("CUDA_VISIBLE_DEVICES") is None:
-    os.environ["CUDA_VISIBLE_DEVICES"] = "7"
+# if os.getenv("CUDA_VISIBLE_DEVICES") is None:
+    # os.environ["CUDA_VISIBLE_DEVICES"] = "7"