cpu or cuda

app.py

@@ -230,6 +230,9 @@ def calculate_iou(box1, box2):
 
 # @spaces.GPU(enable_queue=True, duration=120)
 def buildmodel(**kwargs):
+    global model
+    global quantizer
+    global tokenizer
     from modeling_crello import CrelloModel, CrelloModelConfig
     from quantizer import get_quantizer
     # seed / input model / resume
@@ -289,7 +292,6 @@ def buildmodel(**kwargs):
         config=model_args,
         # cache_dir="/openseg_blob/v-yanbin/GradioDemo/cache_dir",
     ).to("cuda")
-    # model = CrelloModel(config=model_args)
     
     tokenizer.add_special_tokens({"mask_token": "<mask>"}) 
     quantizer.additional_special_tokens.add("<mask>")
@@ -298,7 +300,7 @@ def buildmodel(**kwargs):
     for token in added_special_tokens_list:
         quantizer.additional_special_tokens.add(token)
         
-    return model, quantizer, tokenizer
+    # return model, quantizer, tokenizer
 
 def construction_layout():
     params_dict = {  
@@ -318,7 +320,8 @@ def construction_layout():
     }  
     device = "cuda"
     # Init model
-    model, quantizer, tokenizer = buildmodel(**params_dict)
+    buildmodel(**params_dict)
+    # model, quantizer, tokenizer = buildmodel(**params_dict)
     
     # print('resize token embeddings to match the tokenizer', 129423)
     # model.lm.resize_token_embeddings(129423)
@@ -332,20 +335,24 @@ def construction_layout():
     model.eval()
     # quantizer = quantizer.to("cuda")
     # tokenizer = tokenizer.to("cuda")
-    model.lm = model.lm.to("cuda")
+    # model.lm = model.lm.to("cuda")
     print(model.lm.device)
-    return model, quantizer, tokenizer, params_dict["width"], params_dict["height"], device
+    
+    return params_dict["width"], params_dict["height"], device
+    # return model, quantizer, tokenizer, params_dict["width"], params_dict["height"], device
 
 @torch.no_grad()   
 @spaces.GPU(duration=120) 
 def evaluate_v1(inputs, model, quantizer, tokenizer, width, height, device, do_sample=False, temperature=1.0, top_p=1.0, top_k=50):
+    print(model.lm.device)
     json_example = inputs
     input_intension = '{"wholecaption":"' + json_example["wholecaption"] + '","layout":[{"layer":'
     print("tokenizer1")
     inputs = tokenizer(
         input_intension, return_tensors="pt"
     ).to(model.lm.device)
-    # print(inputs.device)
+    print("Input IDs device:", inputs["input_ids"].device)
+    print("Attention Mask device:", inputs["attention_mask"].device)
     print("tokenizer2")
     
     stopping_criteria = StoppingCriteriaList()
@@ -379,12 +386,6 @@ def evaluate_v1(inputs, model, quantizer, tokenizer, width, height, device, do_s
     return pred_json_example
 
 def inference(generate_method, intention, model, quantizer, tokenizer, width, height, device, do_sample=True, temperature=1.0, top_p=1.0, top_k=50):
-    # def FormulateInput(intension: str):
-    #     resdict = {}
-    #     resdict["wholecaption"] = intension
-    #     resdict["layout"] = []
-    #     return resdict
-    # rawdata = FormulateInput(intention)
     rawdata = {}
     rawdata["wholecaption"] = intention
     rawdata["layout"] = []
@@ -525,7 +526,8 @@ def process_svg(text_input, tuple_input, seed, true_gs, inference_steps):
     return result_images, svg_file_path, svg_editor
     
 def main():
-    model, quantizer, tokenizer, width, height, device = construction_layout()
+    # model, quantizer, tokenizer, width, height, device = construction_layout()
+    width, height, device = construction_layout()
     
     inference_partial = partial(
         inference,
@@ -540,12 +542,17 @@ def main():
     def process_preddate(intention, temperature, top_p, generate_method='v1'):
         intention = intention.replace('\n', '').replace('\r', '').replace('\\', '')
         intention = ensure_space_after_period(intention)
+        print(f"process_preddate: {model.lm.device}")
+        model.lm.to("cuda")
+        print(f"after process_preddate: {model.lm.device}")
         if temperature == 0.0:
             # print("looking for greedy decoding strategies, set `do_sample=False`.")
-            preddata = inference_partial(generate_method, intention, do_sample=False)
+            # preddata = inference_partial(generate_method, intention, do_sample=False)
+            preddata = inference(generate_method, intention, model=model, quantizer=quantizer, tokenizer=tokenizer, width=width, height=height, device=device, do_sample=False)
         else:
-            preddata = inference_partial(generate_method, intention, temperature=temperature, top_p=top_p)
-        # wholecaption = preddata["wholecaption"]
+            # preddata = inference_partial(generate_method, intention, temperature=temperature, top_p=top_p)
+            preddata = inference(generate_method, intention, model=model, quantizer=quantizer, tokenizer=tokenizer, width=width, height=height, device=device, temperature=temperature, top_p=top_p)
+            
         layouts = preddata["layout"]
         list_box = []
         for i, layout in enumerate(layouts):