Merge pull request #40 from fa0311/add-sub-folder

add subfolder option

frontend/modules/api/models/BuildRequest.ts

@@ -25,6 +25,12 @@ export interface BuildRequest {
      * @memberof BuildRequest
      */
     model_id: string;
+    /**
+     * 
+     * @type {string}
+     * @memberof BuildRequest
+     */
+    subfolder?: string;
     /**
      * 
      * @type {string}
@@ -132,6 +138,7 @@ export function BuildRequestFromJSONTyped(json: any, ignoreDiscriminator: boolea
     return {
 
         'model_id': json['model_id'],
+        'subfolder': !exists(json, 'subfolder') ? undefined : json['subfolder'],
         'hf_token': !exists(json, 'hf_token') ? undefined : json['hf_token'],
         'fp16': !exists(json, 'fp16') ? undefined : json['fp16'],
         'verbose': !exists(json, 'verbose') ? undefined : json['verbose'],
@@ -159,6 +166,7 @@ export function BuildRequestToJSON(value?: BuildRequest | null): any {
     return {
 
         'model_id': value.model_id,
+        'subfolder': value.subfolder,
         'hf_token': value.hf_token,
         'fp16': value.fp16,
         'verbose': value.verbose,

frontend/src/tabs/engine.tsx

@@ -89,6 +89,18 @@ const Engine = () => {
                 onChange={(e) => setForm({ ...form, model_id: e.currentTarget.value })}
               />
             </Input.Wrapper>
+            <Input.Wrapper label={'Sub Folder'}>
+              <Input
+                placeholder=""
+                defaultValue={form.hf_token}
+                onChange={(e) =>
+                  setForm({
+                    ...form,
+                    subfolder: e.currentTarget.value,
+                  })
+                }
+              />
+            </Input.Wrapper>
 
             <Input.Wrapper label={'Hugging Face Access Token'}>
               <Input

modules/api/routes/engine_api.py

@@ -11,6 +11,7 @@
 
 class BuildRequest(BaseModel):
     model_id: str
+    subfolder:str = ""
     hf_token: str = ""
     fp16: bool = False
     verbose: bool = False
@@ -49,6 +50,7 @@ async def build_engine(req: BuildRequest):
 
     builder = EngineBuilder(
         model_id=req.model_id,
+        subfolder=req.subfolder,
         hf_token=req.hf_token,
         fp16=req.fp16,
         verbose=req.verbose,

modules/diffusion/tensorrt/engine.py

@@ -86,6 +86,7 @@ class EngineBuilder:
     def __init__(
         self,
         model_id: str,
+        subfolder: str="",
         hf_token="",
         fp16=False,
         verbose=False,
@@ -102,6 +103,8 @@ def __init__(
         onnx_minimal_optimization=False,
     ):
         self.device = "cuda"
+        self.model_id = model_id
+        self.subfolder = subfolder
         self.hf_token = hf_token
         self.fp16 = fp16
         self.verbose = verbose
@@ -118,6 +121,7 @@ def __init__(
         self.models = {
             "unet": UNet(
                 model_id,
+                subfolder=subfolder,
                 hf_token=hf_token,
                 fp16=fp16,
                 device=self.device,
@@ -126,6 +130,7 @@ def __init__(
             ),
             "vae": VAE(
                 model_id,
+                subfolder=subfolder,
                 hf_token=hf_token,
                 device=self.device,
                 verbose=verbose,
@@ -136,6 +141,7 @@ def __init__(
         self.model_dir = os.path.join(
             config.get("model_dir"),
             os.path.basename(model_id) if os.path.isabs(model_id) else model_id,
+            subfolder,
         )
 
     def build(self, generator=False, on_end=lambda: ()):
@@ -180,6 +186,8 @@ def build(self, generator=False, on_end=lambda: ()):
             "onnx_opset": self.onnx_opset,
             "build_static_batch": self.build_static_batch,
             "build_dynamic_shape": self.build_dynamic_shape,
+            "model_id":self.model_id,
+            "subfolder":self.subfolder,
         }
         txt = json.dumps(meta)
         with open(os.path.join(self.model_dir, "model_index.json"), mode="w") as f:

modules/diffusion/tensorrt/models.py

@@ -1,4 +1,5 @@
 import torch
+import os
 from transformers import CLIPTextModel
 
 from lib.trt import models
@@ -33,6 +34,7 @@ class UNet(models.UNet):
     def __init__(
         self,
         model_id: str,
+        subfolder:str ="",
         hf_token="",
         text_maxlen=77,
         embedding_dim=768,
@@ -45,20 +47,25 @@ def __init__(
             hf_token, text_maxlen, embedding_dim, fp16, device, verbose, max_batch_size
         )
         self.model_id = model_id
+        self.subfolder = subfolder
 
     def get_model(self):
         model_opts = (
             {"revision": "fp16", "torch_dtype": torch.float16} if self.fp16 else {}
         )
         return UNet2DConditionModel.from_pretrained(
-            self.model_id, subfolder="unet", use_auth_token=self.hf_token, **model_opts
+            self.model_id,
+            subfolder=os.path.join(self.subfolder, "unet").replace(os.sep, "/"),
+            use_auth_token=self.hf_token,
+            **model_opts,
         ).to(self.device)
 
 
 class VAE(models.VAE):
     def __init__(
         self,
         model_id: str,
+        subfolder:str ="",
         hf_token="",
         text_maxlen=77,
         embedding_dim=768,
@@ -71,11 +78,12 @@ def __init__(
             hf_token, text_maxlen, embedding_dim, fp16, device, verbose, max_batch_size
         )
         self.model_id = model_id
+        self.subfolder = subfolder
 
     def get_model(self):
         vae = AutoencoderKL.from_pretrained(
             self.model_id,
-            subfolder="vae",
+            subfolder=os.path.join(self.subfolder, "vae").replace(os.sep, "/"),
             use_auth_token=self.hf_token,
         ).to(self.device)
         vae.forward = vae.decode

modules/diffusion/tensorrt/runner.py

@@ -161,9 +161,9 @@ def __init__(self, model_dir: str):
                 "openai/clip-vit-large-patch14", fp16=self.fp16, device=self.device
             ),
             "unet": UNet(
-                self.meta["models"]["unet"], fp16=self.fp16, device=self.device
+                self.meta["models"]["unet"], subfolder=self.meta["subfolder"], fp16=self.fp16, device=self.device
             ),
-            "vae": VAE(self.meta["models"]["vae"], fp16=self.fp16, device=self.device),
+            "vae": VAE(self.meta["models"]["vae"], subfolder=self.meta["subfolder"], fp16=self.fp16, device=self.device),
         }
         self.en_vae = self.models["vae"].get_model()
 

modules/runners.py

@@ -25,6 +25,8 @@ def set_runner(
             meta = json.loads(f.read())
         if "model_id" not in meta:
             meta["model_id"] = os.path.relpath(model_dir).replace(os.sep, "/")
+        if "subfolder" not in meta:
+            meta["subfolder"] = ""
         with open(meta_path, mode="w") as f:
             f.write(json.dumps(meta))
     except Exception as e: