Merge pull request #1 from abesmon/hf-key

custom hugging face models

gradio_app.py

@@ -67,6 +67,8 @@
 parser.add_argument('--light_phi', type=float, default=0, help="default GUI light direction in [0, 360), azimuth")
 parser.add_argument('--max_spp', type=int, default=1, help="GUI rendering max sample per pixel")
 
+parser.add_argument('--hf_key', type=str, default=None, help="hugging face Stable diffusion model key")
+
 opt = parser.parse_args() 
 
 # default to use -O !!!
@@ -91,7 +93,7 @@
 
 if opt.guidance == 'stable-diffusion':
     from nerf.sd import StableDiffusion
-    guidance = StableDiffusion(device)
+    guidance = StableDiffusion(device, opt.hf_key)
 elif opt.guidance == 'clip':
     from nerf.clip import CLIP
     guidance = CLIP(device)

nerf/sd.py

@@ -17,15 +17,17 @@ def seed_everything(seed):
     #torch.backends.cudnn.benchmark = True
 
 class StableDiffusion(nn.Module):
-    def __init__(self, device, sd_version='2.0'):
+    def __init__(self, device, sd_version='2.0', hf_key=None):
         super().__init__()
 
         self.device = device
         self.sd_version = sd_version
 
         print(f'[INFO] loading stable diffusion...')
 
-        if self.sd_version == '2.0':
+        if hf_key:
+            model_key = hf_key
+        elifif self.sd_version == '2.0':
             model_key = "stabilityai/stable-diffusion-2-base"
         elif self.sd_version == '1.5':
             model_key = "runwayml/stable-diffusion-v1-5"