use safetensors

app.py

@@ -228,7 +228,7 @@ def calculate_iou(box1, box2):
     iou = intersection_area / union_area
     return iou
 
-# @spaces.GPU(enable_queue=True, duration=60)
+# @spaces.GPU(enable_queue=True, duration=120)
 def buildmodel(**kwargs):
     from modeling_crello import CrelloModel, CrelloModelConfig
     from quantizer import get_quantizer
@@ -337,7 +337,7 @@ def construction_layout():
     return model, quantizer, tokenizer, params_dict["width"], params_dict["height"], device
 
 @torch.no_grad()   
-@spaces.GPU(duration=60) 
+@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):
     json_example = inputs
     input_intension = '{"wholecaption":"' + json_example["wholecaption"] + '","layout":[{"layer":'
@@ -379,13 +379,15 @@ 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)
+    # def FormulateInput(intension: str):
+    #     resdict = {}
+    #     resdict["wholecaption"] = intension
+    #     resdict["layout"] = []
+    #     return resdict
+    # rawdata = FormulateInput(intention)
+    rawdata = {}
+    rawdata["wholecaption"] = intention
+    rawdata["layout"] = []
     
     if generate_method == 'v1':
         max_try_time = 5
@@ -399,7 +401,7 @@ def inference(generate_method, intention, model, quantizer, tokenizer, width, he
 
     return preddata
 
-# @spaces.GPU(enable_queue=True, duration=60)
+# @spaces.GPU(enable_queue=True, duration=120)
 def construction():
     global pipeline
     global transp_vae
@@ -418,6 +420,7 @@ def construction():
         "WYBar/ART_test_weights",
         subfolder="custom_vae", 
         torch_dtype=torch.float32,
+        use_safetensors=True,
         # cache_dir="/openseg_blob/v-yanbin/GradioDemo/cache_dir"
     )
     
@@ -433,7 +436,7 @@ def construction():
     
     # return pipeline, transp_vae
 
-@spaces.GPU(duration=60)
+@spaces.GPU(duration=120)
 def test_one_sample(validation_box, validation_prompt, true_gs, inference_steps, pipeline, generator, transp_vae):
     print(validation_box)
     output, rgba_output, _, _ = pipeline(