fix crello pretrained

app.py

@@ -50,10 +50,11 @@ def generate_unique_filename():
     unique_filename = f"{timestamp}"
     return unique_filename
 
+git_token = os.environ.get("GIT_TOKEN")
 def upload_to_github(file_path, 
                      repo='WYBar/gradiodemo_svg', 
                      branch='main', 
-                     token='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'):
+                     token=git_token):
     if not os.path.isfile(file_path):
         print(f"File not found: {file_path}")
         return None
@@ -274,26 +275,21 @@ def buildmodel(**kwargs):
         pad_token_id=tokenizer.pad_token_id,
         ignore_ids=tokenizer.convert_tokens_to_ids(quantizer.ignore_tokens), 
     )
-    model_args.freeze_lm = True
-    model_args.opt_version = "WYBar/LLM_For_Layout_Planning"
+    model_args.freeze_lm = False
+    model_args.opt_version = input_model
     model_args.use_lora = False
     model_args.load_in_4bit = kwargs.get('load_in_4bit', False)
     # model = CrelloModel.from_pretrained(
     #     resume,
     #     config=model_args
     # ).to(device)
-    # model = CrelloModel.from_pretrained(
-    #     "WYBar/LLM_For_Layout_Planning",
-    #     subfolder="checkpoint-26000",      # 加载检查点目录
-    #     config=model_args,
-    #     # cache_dir="/openseg_blob/v-yanbin/GradioDemo/cache_dir",
-    # )
-    model = CrelloModel(config=model_args)
-    print("before .to(device)")
-    model = model.to(device)
-    print("after .to(device)")
-    model = model.bfloat16()
-    model.eval()
+    model = CrelloModel.from_pretrained(
+        "WYBar/LLM_For_Layout_Planning",
+        subfolder="checkpoint-26000",      # 加载检查点目录
+        config=model_args,
+        # cache_dir="/openseg_blob/v-yanbin/GradioDemo/cache_dir",
+    )
+    # model = CrelloModel(config=model_args)
     
     tokenizer.add_special_tokens({"mask_token": "<mask>"}) 
     quantizer.additional_special_tokens.add("<mask>")
@@ -328,6 +324,12 @@ def construction_layout():
     model.lm.resize_token_embeddings(129423)
     model.input_embeddings = model.lm.get_input_embeddings()
     print('after token embeddings to match the tokenizer', 129423)
+    
+    print("before .to(device)")
+    model = model.to(device)
+    print("after .to(device)")
+    model = model.bfloat16()
+    model.eval()
     return model, quantizer, tokenizer, params_dict["width"], params_dict["height"], device
 
 @torch.no_grad()   

app_test.py

@@ -50,10 +50,11 @@ def generate_unique_filename():
     unique_filename = f"{timestamp}"
     return unique_filename
 
+git_token = os.environ.get("GIT_TOKEN")
 def upload_to_github(file_path, 
                      repo='WYBar/gradiodemo_svg', 
                      branch='main', 
-                     token='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'):
+                     token=git_token):
     if not os.path.isfile(file_path):
         print(f"File not found: {file_path}")
         return None
@@ -274,26 +275,22 @@ def buildmodel(**kwargs):
         pad_token_id=tokenizer.pad_token_id,
         ignore_ids=tokenizer.convert_tokens_to_ids(quantizer.ignore_tokens), 
     )
-    model_args.freeze_lm = True
-    model_args.opt_version = "WYBar/LLM_For_Layout_Planning"
+    model_args.freeze_lm = False
+    model_args.opt_version = input_model
     model_args.use_lora = False
     model_args.load_in_4bit = kwargs.get('load_in_4bit', False)
     # model = CrelloModel.from_pretrained(
     #     resume,
     #     config=model_args
     # ).to(device)
-    # model = CrelloModel.from_pretrained(
-    #     "WYBar/LLM_For_Layout_Planning",
-    #     subfolder="checkpoint-26000",      # 加载检查点目录
-    #     config=model_args,
-    #     # cache_dir="/openseg_blob/v-yanbin/GradioDemo/cache_dir",
-    # )
-    model = CrelloModel(config=model_args)
-    print("before .to(device)")
-    model = model.to(device)
-    print("after .to(device)")
-    model = model.bfloat16()
-    model.eval()
+
+    model = CrelloModel.from_pretrained(
+        "WYBar/LLM_For_Layout_Planning",
+        subfolder="checkpoint-26000",      # 加载检查点目录
+        config=model_args,
+        cache_dir="/openseg_blob/v-yanbin/GradioDemo/cache_dir",
+    )
+    # model = CrelloModel(config=model_args)
     
     tokenizer.add_special_tokens({"mask_token": "<mask>"}) 
     quantizer.additional_special_tokens.add("<mask>")
@@ -307,8 +304,8 @@ def buildmodel(**kwargs):
 def construction_layout():
     params_dict = {  
         # 需要修改
-        "input_model": "WYBar/LLM_For_Layout_Planning", 
-        "resume": "WYBar/LLM_For_Layout_Planning",
+        "input_model": "/openseg_blob/v-sirui/temporary/2024-02-21/Layout_train/COLEv2/Design_LLM/checkpoint/Meta-Llama-3-8B", 
+        "resume": "/openseg_blob/v-sirui/temporary/2024-02-21/SVD/Int2lay_1016/checkpoint/int2lay_1031/1031_test/checkpoint-26000/",
         
         "seed": 0,  
         "mask_values": False,  
@@ -328,6 +325,12 @@ def construction_layout():
     model.lm.resize_token_embeddings(129423)
     model.input_embeddings = model.lm.get_input_embeddings()
     print('after token embeddings to match the tokenizer', 129423)
+    
+    print("before .to(device)")
+    model = model.to(device)
+    print("after .to(device)")
+    model = model.bfloat16()
+    model.eval()
     return model, quantizer, tokenizer, params_dict["width"], params_dict["height"], device
 
 @torch.no_grad()   
@@ -678,7 +681,7 @@ def main():
             inputs=[intention_input, temperature_input, top_p_input, seed_input, true_gs_input, inference_steps_input], 
             outputs=[list_box_output, result_images, svg_file, svg_editor, text_input, tuple_input]
         )
-    demo.launch()
+    demo.launch(server_name='0.0.0.0', server_port=7860)
 
 if __name__ == "__main__":
     main()

modeling_crello.py

@@ -1,6 +1,5 @@
 import torch
-from transformers import PreTrainedModel, PretrainedConfig, AutoModel, AutoModelForCausalLM, OPTForCausalLM
-# from transformers import BitsAndBytesConfig
+from transformers import PreTrainedModel, PretrainedConfig, AutoModel, AutoModelForCausalLM, OPTForCausalLM, BitsAndBytesConfig
 from torch import nn
 import os
 from typing import Optional, List
@@ -117,12 +116,13 @@ class CrelloModel(PreTrainedModel):
 
   def __init__(self, config: CrelloModelConfig): # 显示声明config类型
     super().__init__(config)
+    use_auth_token = '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'
 
     self.pad_token_id = config.pad_token_id
 
     self.args = config
 
-    opt_version = "WYBar/LLM_For_Layout_Planning"
+    opt_version = config.opt_version
 
     print(f"Using {opt_version} for the language model.")
 
@@ -132,7 +132,9 @@ class CrelloModel(PreTrainedModel):
     else:
       if config.load_in_4bit:
         print("\n would load_in_4bit")
-        quantization_config = None
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=config.load_in_4bit
+        )
         # This means: fit the entire model on the GPU:0
         local_rank = int(os.environ.get("LOCAL_RANK", 0))
         device_map = {"": local_rank}
@@ -151,8 +153,20 @@ class CrelloModel(PreTrainedModel):
         # device_map=device_map,
         trust_remote_code=True,
         torch_dtype=torch.bfloat16,
-        # cache_dir="/openseg_blob/v-yanbin/GradioDemo/cache_dir",
+        cache_dir="/openseg_blob/v-yanbin/GradioDemo/cache_dir",
       )
+      # self.lm = AutoModelForCausalLM.from_pretrained(
+      #   opt_version,
+      #   use_auth_token=use_auth_token,
+      #   quantization_config=quantization_config,
+      #   device_map=device_map,
+      #   trust_remote_code=True,
+      #   # attn_implementation="flash_attention_2",
+      #   # flash_attn=True,
+      #   # flash_rotary=True,
+      #   # fused_dense=True,
+      #   torch_dtype=torch.bfloat16,
+      # )
       word_embed_proj_dim = self.lm.config.hidden_size
     self.config.hidden_size = self.lm.config.hidden_size
     self.opt_version = opt_version
@@ -160,8 +174,8 @@ class CrelloModel(PreTrainedModel):
     if self.args.freeze_lm:
       self.lm.eval()
       print("Freezing the LM.")
-      # for param in self.lm.parameters():
-      #   param.requires_grad = False
+      for param in self.lm.parameters():
+        param.requires_grad = False
     else:
       print("\n no freeze lm, so to train lm")
       self.lm.train()
@@ -170,7 +184,6 @@ class CrelloModel(PreTrainedModel):
     # print('resize token embeddings to match the tokenizer', config.vocab_size)
     # self.lm.resize_token_embeddings(config.vocab_size)
     # self.input_embeddings = self.lm.get_input_embeddings()
-    # print('after token embeddings to match the tokenizer', config.vocab_size)
     
   def train(self, mode=True):
     super().train(mode=mode)