add some more prov imgs

app.py

@@ -13,7 +13,7 @@ from demo.rm_bg import rm_bg
 
 with gr.Blocks(title="DiLightNet Demo") as demo:
     gr.Markdown("""# DiLightNet: Fine-grained Lighting Control for Diffusion-based Image Generation
-                ## A demo for generating images under point/environmantal lighting using DiLightNet. For full usage (video generation & arbitary lighting condition), please refer to our [GitHub repository](https://github.com/iamNCJ/DiLightNet)""")
+                ## A demo for generating images under point/environmantal lighting using DiLightNet. For full usage (video generation & arbitary lighting condition & depth-conditioned generation) and more examples, please refer to our [GitHub repository](https://github.com/iamNCJ/DiLightNet)""")
 
     with gr.Row():
         # 1. Reference Image Input / Generation
@@ -34,7 +34,7 @@ with gr.Blocks(title="DiLightNet Demo") as demo:
             gr.Examples(
                 examples=[os.path.join("examples/provisional_img", i) for i in os.listdir("examples/provisional_img")],
                 inputs=[input_image],
-                examples_per_page = 20,
+                examples_per_page=8,
             )
 
         # 2. Background Removal
@@ -81,7 +81,7 @@ with gr.Blocks(title="DiLightNet Demo") as demo:
                 env_examples = gr.Examples(
                     examples=[[os.path.join("examples/env_map_preview", i), os.path.join("examples/env_map", i).replace("png", "exr")] for i in os.listdir("examples/env_map_preview")],
                     inputs=[env_map_preview, env_map_path],
-                    examples_per_page = 20,
+                    examples_per_page=20,
                 )
                 render_btn_env = gr.Button(value="Render Hints")
 
@@ -164,13 +164,15 @@ with gr.Blocks(title="DiLightNet Demo") as demo:
                     seed=int(relighting_seed),
                     cfg=relighting_cfg
                 )
-                relit_img = imageio.v3.imread(res_folder_path + '/relighting00.png')
+                relit_img = imageio.v3.imread(res_folder_path + '/relighting00_0.png')
                 if do_env_inpainting:
                     bg = imageio.v3.imread(res_folder_path + f'/bg00.png') / 255.
-                    relit_img = relit_img / 255.
-                    mask_for_bg = imageio.v3.imread(res_folder_path + '/hint00_diffuse.png')[..., -1:] / 255.
-                    relit_img = relit_img * mask_for_bg + bg * (1. - mask_for_bg)
-                    relit_img = (relit_img * 255).clip(0, 255).astype(np.uint8)
+                else:
+                    bg = np.zeros_like(relit_img)
+                relit_img = relit_img / 255.
+                mask_for_bg = imageio.v3.imread(res_folder_path + '/hint00_diffuse.png')[..., -1:] / 255.
+                relit_img = relit_img * mask_for_bg + bg * (1. - mask_for_bg)
+                relit_img = (relit_img * 255).clip(0, 255).astype(np.uint8)
                 return relit_img
 
             relighting_generate_btn.click(fn=gen_relighting_image,

demo/relighting_gen.py

@@ -1,8 +1,7 @@
 import imageio
 import numpy as np
-import spaces
 import torch
-from diffusers import UniPCMultistepScheduler, StableDiffusionControlNetPipeline, StableDiffusionInpaintPipeline, ConsistencyDecoderVAE
+from diffusers import UniPCMultistepScheduler, StableDiffusionControlNetPipeline, StableDiffusionInpaintPipeline
 from diffusers.utils import get_class_from_dynamic_module
 
 from tqdm import tqdm
@@ -22,10 +21,8 @@ controlnet = NeuralTextureControlNetModel.from_pretrained(
     "dilightnet/DiLightNet-submissions-300k",
     torch_dtype=dtype,
 )
-vae = ConsistencyDecoderVAE.from_pretrained("openai/consistency-decoder", torch_dtype=dtype)
 pipe = StableDiffusionControlNetPipeline.from_pretrained(
     "stabilityai/stable-diffusion-2-1",
-    vae=vae,
     controlnet=controlnet,
     torch_dtype=dtype
 ).to(device)
@@ -40,7 +37,7 @@ inpainting_pipe.set_progress_bar_config(disable=True)
 
 
 @spaces.GPU
-def relighting_gen(masked_ref_img, mask, cond_path, frames, prompt, steps, seed, cfg, inpaint=False):
+def relighting_gen(masked_ref_img, mask, cond_path, frames, prompt, steps, seed, cfg, num_imgs_per_prompt=1, inpaint=False):
     mask = mask[..., :1] / 255.
     for i in tqdm(range(frames)):
         source_image = masked_ref_img[..., :3] / 255.
@@ -55,13 +52,13 @@ def relighting_gen(masked_ref_img, mask, cond_path, frames, prompt, steps, seed,
             images.append(image)
 
         hint = np.concatenate(images, axis=2).astype(np.float32)[None]
-
         hint = torch.from_numpy(hint).to(dtype).permute(0, 3, 1, 2).to(device)
-        generator = torch.manual_seed(seed)
-        image = pipe(
-            prompt, num_inference_steps=steps, generator=generator, image=hint, num_images_per_prompt=1, guidance_scale=cfg, output_type='np',
-        ).images[0]  # [H, W, C]
+        generator = torch.Generator(device=device).manual_seed(seed)
+        images = pipe(
+            prompt, num_inference_steps=steps, generator=generator, image=hint, num_images_per_prompt=num_imgs_per_prompt, guidance_scale=cfg, output_type='np',
+        ).images  # [N, H, W, C]
         if inpaint:
             mask_image = (1. - mask)[None]
-            image = inpainting_pipe(prompt=prompt, image=image[None], mask_image=mask_image, generator=generator, output_type='np', cfg=3.0, strength=1.0).images[0]
-        imageio.imwrite(f'{cond_path}/relighting{i:02d}.png', (image * 255).clip(0, 255).astype(np.uint8))
+            images = inpainting_pipe(prompt=prompt, image=images, mask_image=mask_image, generator=generator, output_type='np', cfg=3.0, strength=1.0).images
+        for idx in range(num_imgs_per_prompt):
+            imageio.imwrite(f'{cond_path}/relighting{i:02d}_{idx}.png', (images[idx] * 255).clip(0, 255).astype(np.uint8))

demo/render_hints.py

@@ -27,6 +27,8 @@ def render_hint_images(model_path, fov, pls, power=500., geo_smooth=True, output
             bpy.context.preferences.addons["cycles"].preferences.get_devices()
             bpy.context.scene.cycles.device = 'GPU'
             bpy.context.preferences.addons['cycles'].preferences.compute_device_type = 'CUDA'
+            bpy.context.scene.render.threads = 8
+            bpy.context.scene.render.threads_mode = 'FIXED'
 
         # Enable the alpha channel for GT mask
         bpy.context.scene.render.film_transparent = True
@@ -85,7 +87,7 @@ def render_hint_images(model_path, fov, pls, power=500., geo_smooth=True, output
     return output_folder
 
 
-def render_bg_images(fov, pls, output_folder: Optional[str] = None, env_map: Optional[str] = None, env_start_azi=0., resolution=512):
+def render_bg_images(fov, pls, output_folder: Optional[str] = None, env_map: Optional[str] = None, env_start_azi=0., resolution=512, use_gpu=False):
     import bpy
     import numpy as np
 
@@ -105,6 +107,13 @@ def render_bg_images(fov, pls, output_folder: Optional[str] = None, env_map: Opt
         bpy.context.scene.render.film_transparent = False
         bpy.context.scene.render.image_settings.color_mode = 'RGB'
 
+        if use_gpu:
+            bpy.context.preferences.addons["cycles"].preferences.get_devices()
+            bpy.context.scene.cycles.device = 'GPU'
+            bpy.context.preferences.addons['cycles'].preferences.compute_device_type = 'CUDA'
+            bpy.context.scene.render.threads = 8
+            bpy.context.scene.render.threads_mode = 'FIXED'
+
     def render_env_bg(output_path):
         bpy.context.scene.view_layers["ViewLayer"].material_override = None
         bpy.context.scene.render.image_settings.file_format = 'PNG'