v1

app.py

@@ -62,60 +62,78 @@ def bot_streaming(message, history, link, temperature, new_max_token, top_p):
     if "1.8B" in link:
         model = model_1_8
         tokenizer = tokenizer_1_8
+        path = "BK-Lee/TroL-1.8B"
     elif "3.8B" in link:
         model = model_3_8
         tokenizer = tokenizer_3_8
+        path = "BK-Lee/TroL-3.8B"
     elif "7B" in link:
         model = model_7
         tokenizer = tokenizer_7
+        path = "BK-Lee/TroL-7B"
     
+    # trol gating load
+    from huggingface_hub import hf_hub_download
+    try:
+        model.model.initialize_trol_gating()
+        model.model.trol_gating.load_state_dict(torch.load(hf_hub_download(repo_id=path, filename="trol_gating.pt")))
+    except:
+        model.language_model.model.initialize_trol_gating()
+        model.language_model.model.trol_gating.load_state_dict(torch.load(hf_hub_download(repo_id=path, filename="trol_gating.pt")))
+
+    # X -> float16 conversion 
+    for param in model.parameters():
+        if 'float32' in str(param.dtype).lower() or 'float16' in str(param.dtype).lower():
+            param.data = param.data.to(torch.float16)
+
     # cpu -> gpu
     for param in model.parameters():
         if not param.is_cuda:
             param.data = param.to(accel.device)
 
-    # prompt type -> input prompt
-    image_token_number = None
-    if len(message['files']) == 1:
-        # Image Load
-        image = pil_to_tensor(Image.open(message['files'][0]).convert("RGB"))
-        if "3.8B" not in link:
-            image_token_number = 1225
-            image = F.interpolate(image.unsqueeze(0), size=(490, 490), mode='bicubic').squeeze(0)
-        inputs = [{'image': image.to(accel.device), 'question': message['text']}]
-    elif len(message['files']) > 1:
-        raise Exception("No way!")
-    else:
-        inputs = [{'question': message['text']}]
-
-    # Text Generation
-    with torch.inference_mode():
-        # kwargs
-        streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
-
-        # Threading generation
-        thread = Thread(target=threading_function, kwargs=dict(inputs=inputs,
-                                                            image_token_number=image_token_number,
-                                                            streamer=streamer,
-                                                            model=model,
-                                                            tokenizer=tokenizer,
-                                                            device=accel.device,
-                                                            temperature=temperature,
-                                                            new_max_token=new_max_token,
-                                                            top_p=top_p))
-        thread.start()
-
-        # generated text
-        generated_text = ""
-        for new_text in streamer:
-            generated_text += new_text
-        generated_text
-
-    # Text decoding
-    response = output_filtering(generated_text, model)
-
-    # except:
-    #     response = "There may be unsupported format: ex) pdf, video, sound. Only supported is a single image in this version."
+    try:
+        # prompt type -> input prompt
+        image_token_number = None
+        if len(message['files']) == 1:
+            # Image Load
+            image = pil_to_tensor(Image.open(message['files'][0]).convert("RGB"))
+            if "3.8B" not in link:
+                image_token_number = 1225
+                image = F.interpolate(image.unsqueeze(0), size=(490, 490), mode='bicubic').squeeze(0)
+            inputs = [{'image': image.to(accel.device), 'question': message['text']}]
+        elif len(message['files']) > 1:
+            raise Exception("No way!")
+        else:
+            inputs = [{'question': message['text']}]
+
+        # Text Generation
+        with torch.inference_mode():
+            # kwargs
+            streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
+
+            # Threading generation
+            thread = Thread(target=threading_function, kwargs=dict(inputs=inputs,
+                                                                image_token_number=image_token_number,
+                                                                streamer=streamer,
+                                                                model=model,
+                                                                tokenizer=tokenizer,
+                                                                device=accel.device,
+                                                                temperature=temperature,
+                                                                new_max_token=new_max_token,
+                                                                top_p=top_p))
+            thread.start()
+
+            # generated text
+            generated_text = ""
+            for new_text in streamer:
+                generated_text += new_text
+            generated_text
+
+        # Text decoding
+        response = output_filtering(generated_text, model)
+
+    except:
+        response = "There may be unsupported format: ex) pdf, video, sound. Only supported is a single image in this version."
 
     # private log print
     text = message['text']

trol/load_trol.py

@@ -80,18 +80,4 @@ def load_trol(link):
 
     # setting config
     setting_trol_config(trol, tok_trol, image_special_token)
-
-    # trol gating load
-    from huggingface_hub import hf_hub_download
-    try:
-        trol.model.initialize_trol_gating()
-        trol.model.trol_gating.load_state_dict(torch.load(hf_hub_download(repo_id=path, filename="trol_gating.pt")))
-    except:
-        trol.language_model.model.initialize_trol_gating()
-        trol.language_model.model.trol_gating.load_state_dict(torch.load(hf_hub_download(repo_id=path, filename="trol_gating.pt")))
-
-    # X -> float16 conversion 
-    for param in trol.parameters():
-        if 'float32' in str(param.dtype).lower() or 'float16' in str(param.dtype).lower():
-            param.data = param.data.to(torch.float16)
     return trol, tok_trol