construction all

app.py

@@ -228,18 +228,36 @@ def calculate_iou(box1, box2):
     iou = intersection_area / union_area
     return iou
 
-# @spaces.GPU(enable_queue=True, duration=120)
-def buildmodel(**kwargs):
+def construction_all():
     global model
     global quantizer
     global tokenizer
+    global pipeline
+    global transp_vae
     from modeling_crello import CrelloModel, CrelloModelConfig
     from quantizer import get_quantizer
+    from custom_model_mmdit import CustomFluxTransformer2DModel
+    from custom_model_transp_vae import AutoencoderKLTransformerTraining as CustomVAE
+    from custom_pipeline import CustomFluxPipelineCfg
+    
+    params_dict = {  
+        "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,  
+        "quantizer_version": 'v4',  
+        "mask_type": 'cm3',  
+        "decimal_quantize_types": [],  
+        "num_mask_tokens": 0,  
+        "width": 512,
+        "height": 512,
+        "device": 0,
+    }  
+    
     # seed / input model / resume
-    resume = kwargs.get('resume', None)
-    seed = kwargs.get('seed', None)
-    input_model = kwargs.get('input_model', None)
-    quantizer_version = kwargs.get('quantizer_version', 'v4')
+    seed = params_dict.get('seed', None)
+    input_model = params_dict.get('input_model', None)
+    quantizer_version = params_dict.get('quantizer_version', 'v4')
     
     set_seed(seed)
     # old_tokenizer = AutoTokenizer.from_pretrained(input_model, trust_remote_code=True)
@@ -261,13 +279,13 @@ def buildmodel(**kwargs):
     quantizer = get_quantizer(
                     quantizer_version, 
                     update_vocab = False,
-                    decimal_quantize_types = kwargs.get('decimal_quantize_types'),
-                    mask_values = kwargs['mask_values'],
-                    width = kwargs['width'],
-                    height = kwargs['height'],
+                    decimal_quantize_types = params_dict.get('decimal_quantize_types'),
+                    mask_values = params_dict['mask_values'],
+                    width = params_dict['width'],
+                    height = params_dict['height'],
                     simplify_json = False,
                     num_mask_tokens = 0, 
-                    mask_type = kwargs.get('mask_type'),
+                    mask_type = params_dict.get('mask_type'),
                 )
     quantizer.setup_tokenizer(tokenizer)  
       
@@ -280,11 +298,7 @@ def buildmodel(**kwargs):
     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_args.load_in_4bit = params_dict.get('load_in_4bit', False)
     
     model = CrelloModel.from_pretrained(
         "WYBar/LLM_For_Layout_Planning",
@@ -300,63 +314,46 @@ def buildmodel(**kwargs):
     for token in added_special_tokens_list:
         quantizer.additional_special_tokens.add(token)
         
-    print(f"before .to(device):{model.device} {model.lm.device}")
+    transformer = CustomFluxTransformer2DModel.from_pretrained(
+        "WYBar/ART_test_weights",
+        subfolder="fused_transformer", 
+        torch_dtype=torch.bfloat16,
+        # cache_dir="/openseg_blob/v-yanbin/GradioDemo/cache_dir"
+    )
+    
+    transp_vae = CustomVAE.from_pretrained(
+        "WYBar/ART_test_weights",
+        subfolder="custom_vae", 
+        torch_dtype=torch.float32,
+        use_safetensors=True,
+        # cache_dir="/openseg_blob/v-yanbin/GradioDemo/cache_dir"
+    )
+    
+    token = os.environ.get("HF_TOKEN")
+    pipeline = CustomFluxPipelineCfg.from_pretrained(
+        "black-forest-labs/FLUX.1-dev",
+        transformer=transformer,
+        torch_dtype=torch.bfloat16,
+        token=token,
+        # cache_dir="/openseg_blob/v-yanbin/GradioDemo/cache_dir"
+    ).to("cuda")
+    pipeline.enable_model_cpu_offload(gpu_id=0) # Save GPU memory
+        
+    print(f"before .to(device):{model.device} {model.lm.device} {pipeline.device}")
     model = model.to("cuda")
-    print(f"after .to(device):{model.device} {model.lm.device}")
+    pipeline = pipeline.to("cuda")
+    print(f"after .to(device):{model.device} {model.lm.device} {pipeline.device}")
     model = model.bfloat16()
     model.eval()
-    # quantizer = quantizer.to("cuda")
-    # tokenizer = tokenizer.to("cuda")
-    # model.lm = model.lm.to("cuda")
-    print(model.lm.device)
-        
-    # return model, quantizer, tokenizer
-
-def construction_layout():
-    params_dict = {  
-        # 需要修改
-        "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,  
-        "quantizer_version": 'v4',  
-        "mask_type": 'cm3',  
-        "decimal_quantize_types": [],  
-        "num_mask_tokens": 0,  
-        "width": 512,
-        "height": 512,
-        "device": 0,
-    }  
-    device = "cuda"
-    # Init model
-    buildmodel(**params_dict)
-    # model, quantizer, tokenizer = buildmodel(**params_dict)
-    
-    # print('resize token embeddings to match the tokenizer', 129423)
-    # 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("cuda")
-    # print("after .to(device)")
-    # model = model.bfloat16()
-    # model.eval()
-    # # quantizer = quantizer.to("cuda")
-    # # tokenizer = tokenizer.to("cuda")
-    # # model.lm = model.lm.to("cuda")
-    # print(model.lm.device)
-    
-    return params_dict["width"], params_dict["height"], device
-    # return model, quantizer, tokenizer, params_dict["width"], params_dict["height"], device
+    print(f"after bf16 & eval .to(device):{model.device} {model.lm.device}  {pipeline.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)
+def evaluate_v1(inputs, model, quantizer, tokenizer, width, height, do_sample=False, temperature=1.0, top_p=1.0, top_k=50):
+    print(f"evaluate_v1 {model.device} {model.lm.device} {pipeline.device}")
     json_example = inputs
     input_intension = '{"wholecaption":"' + json_example["wholecaption"] + '","layout":[{"layer":'
+    
     print("tokenizer1")
     inputs = tokenizer(
         input_intension, return_tensors="pt"
@@ -395,7 +392,7 @@ def evaluate_v1(inputs, model, quantizer, tokenizer, width, height, device, do_s
         pred_json_example = None
     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 inference(generate_method, intention, model, quantizer, tokenizer, width, height, do_sample=True, temperature=1.0, top_p=1.0, top_k=50):
     rawdata = {}
     rawdata["wholecaption"] = intention
     rawdata["layout"] = []
@@ -404,7 +401,7 @@ def inference(generate_method, intention, model, quantizer, tokenizer, width, he
         max_try_time = 5
         preddata = None
         while preddata is None and max_try_time > 0:
-            preddata = evaluate_v1(rawdata, model, quantizer, tokenizer, width, height, device, do_sample=do_sample, temperature=temperature, top_p=top_p, top_k=top_k)
+            preddata = evaluate_v1(rawdata, model, quantizer, tokenizer, width, height, do_sample=do_sample, temperature=temperature, top_p=top_p, top_k=top_k)
             max_try_time -= 1
     else:
         print("Please input correct generate method")
@@ -412,41 +409,6 @@ def inference(generate_method, intention, model, quantizer, tokenizer, width, he
 
     return preddata
 
-# @spaces.GPU(enable_queue=True, duration=120)
-def construction():
-    global pipeline
-    global transp_vae
-    from custom_model_mmdit import CustomFluxTransformer2DModel
-    from custom_model_transp_vae import AutoencoderKLTransformerTraining as CustomVAE
-    from custom_pipeline import CustomFluxPipelineCfg
-
-    transformer = CustomFluxTransformer2DModel.from_pretrained(
-        "WYBar/ART_test_weights",
-        subfolder="fused_transformer", 
-        torch_dtype=torch.bfloat16,
-        # cache_dir="/openseg_blob/v-yanbin/GradioDemo/cache_dir"
-    )
-    
-    transp_vae = CustomVAE.from_pretrained(
-        "WYBar/ART_test_weights",
-        subfolder="custom_vae", 
-        torch_dtype=torch.float32,
-        use_safetensors=True,
-        # cache_dir="/openseg_blob/v-yanbin/GradioDemo/cache_dir"
-    )
-    
-    token = os.environ.get("HF_TOKEN")
-    pipeline = CustomFluxPipelineCfg.from_pretrained(
-        "black-forest-labs/FLUX.1-dev",
-        transformer=transformer,
-        torch_dtype=torch.bfloat16,
-        token=token,
-        # cache_dir="/openseg_blob/v-yanbin/GradioDemo/cache_dir"
-    ).to("cuda")
-    pipeline.enable_model_cpu_offload(gpu_id=0) # Save GPU memory
-    
-    # return pipeline, transp_vae
-
 @spaces.GPU(duration=120)
 def test_one_sample(validation_box, validation_prompt, true_gs, inference_steps, pipeline, generator, transp_vae):
     print(validation_box)
@@ -477,7 +439,7 @@ def test_one_sample(validation_box, validation_prompt, true_gs, inference_steps,
     return output_gradio
 
 def svg_test_one_sample(validation_prompt, validation_box_str, seed, true_gs, inference_steps, pipeline, transp_vae):
-    print("svg_test_one_sample")
+    print(f"svg_test_one_sample {model.device} {model.lm.device} {pipeline.device}")
     generator = torch.Generator().manual_seed(seed)
     try:
         validation_box = ast.literal_eval(validation_box_str)
@@ -511,7 +473,7 @@ def svg_test_one_sample(validation_prompt, validation_box_str, seed, true_gs, in
     return result_images, svg_file_path
 
 def process_svg(text_input, tuple_input, seed, true_gs, inference_steps):
-    print("precess_svg")
+    print(f"precess_svg {model.device} {model.lm.device} {pipeline.device}")
     result_images = []
     result_images, svg_file_path = svg_test_one_sample(text_input, tuple_input, seed, true_gs, inference_steps, pipeline=pipeline, transp_vae=transp_vae)
     # result_images, svg_file_path = gradio_test_one_sample_partial(text_input, tuple_input, seed, true_gs, inference_steps)
@@ -534,64 +496,52 @@ 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()
-    width, height, device = construction_layout()
-    
-    # inference_partial = partial(
-    #     inference,
-    #     model=model,
-    #     quantizer=quantizer,
-    #     tokenizer=tokenizer,
-    #     width=width,
-    #     height=height,
-    #     device=device
-    # )
-    
-    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(generate_method, intention, model=model, quantizer=quantizer, tokenizer=tokenizer, width=width, height=height, device=device, do_sample=False)
+
+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}")
+    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(generate_method, intention, model=model, quantizer=quantizer, tokenizer=tokenizer, width=512, height=512, do_sample=False)
+    else:
+        # preddata = inference_partial(generate_method, intention, temperature=temperature, top_p=top_p)
+        preddata = inference(generate_method, intention, model=model, quantizer=quantizer, tokenizer=tokenizer, width=512, height=512, temperature=temperature, top_p=top_p)
+        
+    layouts = preddata["layout"]
+    list_box = []
+    for i, layout in enumerate(layouts):
+        x, y = layout["x"], layout["y"]
+        width, height = layout["width"], layout["height"]
+        if i == 0:
+            list_box.append((0, 0, width, height))
+            list_box.append((0, 0, width, height))
         else:
-            # 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):
-            x, y = layout["x"], layout["y"]
-            width, height = layout["width"], layout["height"]
-            if i == 0:
-                list_box.append((0, 0, width, height))
-                list_box.append((0, 0, width, height))
-            else:
-                left = x - width // 2
-                top = y - height // 2
-                right = x + width // 2
-                bottom = y + height // 2
-                list_box.append((left, top, right, bottom))
-                
-        # print(list_box)
-        filtered_boxes = list_box[:2]
-        for i in range(2, len(list_box)):
-            keep = True
-            for j in range(1, len(filtered_boxes)):
-                iou = calculate_iou(list_box[i], filtered_boxes[j])
-                if iou > 0.65:
-                    print(list_box[i], filtered_boxes[j])
-                    keep = False
-                    break
-            if keep:
-                filtered_boxes.append(list_box[i])
+            left = x - width // 2
+            top = y - height // 2
+            right = x + width // 2
+            bottom = y + height // 2
+            list_box.append((left, top, right, bottom))
             
-        return str(filtered_boxes), intention, str(filtered_boxes)
+    # print(list_box)
+    filtered_boxes = list_box[:2]
+    for i in range(2, len(list_box)):
+        keep = True
+        for j in range(1, len(filtered_boxes)):
+            iou = calculate_iou(list_box[i], filtered_boxes[j])
+            if iou > 0.65:
+                print(list_box[i], filtered_boxes[j])
+                keep = False
+                break
+        if keep:
+            filtered_boxes.append(list_box[i])
+        
+    return str(filtered_boxes), intention, str(filtered_boxes)
+    
+def main():
+    construction_all()
+    print(f"after construction_all:{model.device} {model.lm.device} {pipeline.device}")
     
     # def process_preddate(intention, generate_method='v1'):
     #     list_box = [(0, 0, 512, 512), (0, 0, 512, 512), (136, 184, 512, 512), (144, 0, 512, 512), (0, 0, 328, 136), (160, 112, 512, 360), (168, 112, 512, 360), (40, 232, 112, 296), (32, 88, 248, 176), (48, 424, 144, 448), (48, 464, 144, 488), (240, 464, 352, 488), (384, 464, 488, 488), (48, 480, 144, 504), (240, 480, 360, 504), (456, 0, 512, 56), (0, 0, 56, 40), (440, 0, 512, 40), (0, 24, 48, 88), (48, 168, 168, 240)]
@@ -600,7 +550,6 @@ def main():
     #     return wholecaption, str(list_box), json_file
 
     # pipeline, transp_vae = construction()
-    construction()
 
     # gradio_test_one_sample_partial = partial(
     #     svg_test_one_sample,