Update app.py

app.py

@@ -178,123 +178,116 @@ global pipeline
 global MultiResNetModel
 global cur_style
 
-@spaces.GPU
-def load_ckpt():
-    global pipeline
-    global MultiResNetModel
-    global cur_style
-    cur_style = 'line + shadow'
-    weight_dtype = torch.float16
+cur_style = 'line + shadow'
+weight_dtype = torch.float16
+
+block_out_channels = [128, 128, 256, 512, 512]
+MultiResNetModel = MultiHiddenResNetModel(block_out_channels, len(block_out_channels))
+MultiResNetModel.load_state_dict(torch.load(os.path.join(model_global_path, 'shadow_GSRP', 'MultiResNetModel.bin'), map_location='cpu'), strict=True)
+MultiResNetModel.to('cuda', dtype=weight_dtype)
 
-    block_out_channels = [128, 128, 256, 512, 512]
-    MultiResNetModel = MultiHiddenResNetModel(block_out_channels, len(block_out_channels))
-    MultiResNetModel.load_state_dict(torch.load(os.path.join(model_global_path, 'shadow_GSRP', 'MultiResNetModel.bin'), map_location='cpu'), strict=True)
-    MultiResNetModel.to('cuda', dtype=weight_dtype)
 
+# transformer
+transform = transforms.Compose([
+            transforms.ToTensor(),  # 将 PIL 图像转换为 Tensor
+            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])  # 归一化
+        ])
+# seed = 43
+lora_rank = 128
+pretrained_model_name_or_path = "PixArt-alpha/PixArt-XL-2-1024-MS"
+
+transformer = PixArtTransformer2DModel.from_pretrained(
+        pretrained_model_name_or_path, subfolder="transformer", revision=None, variant=None
+    )
+pixart_config = get_pixart_config()
+causal_dit = CausalSparseDiTModel(num_attention_heads=pixart_config.get("num_attention_heads"),
+                    attention_head_dim=pixart_config.get("attention_head_dim"),
+                    in_channels=pixart_config.get("in_channels"),
+                    out_channels=pixart_config.get("out_channels"),
+                    num_layers=pixart_config.get("num_layers"),
+                    dropout=pixart_config.get("dropout"),
+                    norm_num_groups=pixart_config.get("norm_num_groups"),
+                    cross_attention_dim=pixart_config.get("cross_attention_dim"),
+                    attention_bias=pixart_config.get("attention_bias"),
+                    sample_size=pixart_config.get("sample_size"),
+                    patch_size=pixart_config.get("patch_size"),
+                    activation_fn=pixart_config.get("activation_fn"),
+                    num_embeds_ada_norm=pixart_config.get("num_embeds_ada_norm"),
+                    upcast_attention=pixart_config.get("upcast_attention"),
+                    norm_type=pixart_config.get("norm_type"),
+                    norm_elementwise_affine=pixart_config.get("norm_elementwise_affine"),
+                    norm_eps=pixart_config.get("norm_eps"),
+                    caption_channels=pixart_config.get("caption_channels"),
+                    attention_type=pixart_config.get("attention_type"))
+
+causal_dit = init_causal_dit(causal_dit, transformer)
+print('loaded causal_dit')
+controlnet = CausalSparseDiTControlModel(num_attention_heads=pixart_config.get("num_attention_heads"),
+                                attention_head_dim=pixart_config.get("attention_head_dim"),
+                                in_channels=pixart_config.get("in_channels"),
+                                cond_chanels = 9,
+                                out_channels=pixart_config.get("out_channels"),
+                                num_layers=pixart_config.get("num_layers"),
+                                dropout=pixart_config.get("dropout"),
+                                norm_num_groups=pixart_config.get("norm_num_groups"),
+                                cross_attention_dim=pixart_config.get("cross_attention_dim"),
+                                attention_bias=pixart_config.get("attention_bias"),
+                                sample_size=pixart_config.get("sample_size"),
+                                patch_size=pixart_config.get("patch_size"),
+                                activation_fn=pixart_config.get("activation_fn"),
+                                num_embeds_ada_norm=pixart_config.get("num_embeds_ada_norm"),
+                                upcast_attention=pixart_config.get("upcast_attention"),
+                                norm_type=pixart_config.get("norm_type"),
+                                norm_elementwise_affine=pixart_config.get("norm_elementwise_affine"),
+                                norm_eps=pixart_config.get("norm_eps"),
+                                caption_channels=pixart_config.get("caption_channels"),
+                                attention_type=pixart_config.get("attention_type")
+                            )
+# controlnet = init_controlnet(controlnet, causal_dit)
+del transformer
+transformer_lora_config = LoraConfig(
+        r=lora_rank,
+        lora_alpha=lora_rank,
+        # use_dora=True,
+        init_lora_weights="gaussian",
+        target_modules=["to_k",
+            "to_q",
+            "to_v",
+            "to_out.0",
+            "proj_in",
+            "proj_out",
+            "ff.net.0.proj",
+            "ff.net.2",
+            "proj",
+            "linear",
+            "linear_1",
+            "linear_2"],#ff.net.0.proj ff.net.2
+    )
+causal_dit.add_adapter(transformer_lora_config)
 
-    # transformer
-    transform = transforms.Compose([
-                transforms.ToTensor(),  # 将 PIL 图像转换为 Tensor
-                transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])  # 归一化
-            ])
-    # seed = 43
-    lora_rank = 128
-    pretrained_model_name_or_path = "PixArt-alpha/PixArt-XL-2-1024-MS"
 
-    transformer = PixArtTransformer2DModel.from_pretrained(
-            pretrained_model_name_or_path, subfolder="transformer", revision=None, variant=None
-        )
-    pixart_config = get_pixart_config()
-    causal_dit = CausalSparseDiTModel(num_attention_heads=pixart_config.get("num_attention_heads"),
-                        attention_head_dim=pixart_config.get("attention_head_dim"),
-                        in_channels=pixart_config.get("in_channels"),
-                        out_channels=pixart_config.get("out_channels"),
-                        num_layers=pixart_config.get("num_layers"),
-                        dropout=pixart_config.get("dropout"),
-                        norm_num_groups=pixart_config.get("norm_num_groups"),
-                        cross_attention_dim=pixart_config.get("cross_attention_dim"),
-                        attention_bias=pixart_config.get("attention_bias"),
-                        sample_size=pixart_config.get("sample_size"),
-                        patch_size=pixart_config.get("patch_size"),
-                        activation_fn=pixart_config.get("activation_fn"),
-                        num_embeds_ada_norm=pixart_config.get("num_embeds_ada_norm"),
-                        upcast_attention=pixart_config.get("upcast_attention"),
-                        norm_type=pixart_config.get("norm_type"),
-                        norm_elementwise_affine=pixart_config.get("norm_elementwise_affine"),
-                        norm_eps=pixart_config.get("norm_eps"),
-                        caption_channels=pixart_config.get("caption_channels"),
-                        attention_type=pixart_config.get("attention_type"))
-
-    causal_dit = init_causal_dit(causal_dit, transformer)
-    print('loaded causal_dit')
-    controlnet = CausalSparseDiTControlModel(num_attention_heads=pixart_config.get("num_attention_heads"),
-                                    attention_head_dim=pixart_config.get("attention_head_dim"),
-                                    in_channels=pixart_config.get("in_channels"),
-                                    cond_chanels = 9,
-                                    out_channels=pixart_config.get("out_channels"),
-                                    num_layers=pixart_config.get("num_layers"),
-                                    dropout=pixart_config.get("dropout"),
-                                    norm_num_groups=pixart_config.get("norm_num_groups"),
-                                    cross_attention_dim=pixart_config.get("cross_attention_dim"),
-                                    attention_bias=pixart_config.get("attention_bias"),
-                                    sample_size=pixart_config.get("sample_size"),
-                                    patch_size=pixart_config.get("patch_size"),
-                                    activation_fn=pixart_config.get("activation_fn"),
-                                    num_embeds_ada_norm=pixart_config.get("num_embeds_ada_norm"),
-                                    upcast_attention=pixart_config.get("upcast_attention"),
-                                    norm_type=pixart_config.get("norm_type"),
-                                    norm_elementwise_affine=pixart_config.get("norm_elementwise_affine"),
-                                    norm_eps=pixart_config.get("norm_eps"),
-                                    caption_channels=pixart_config.get("caption_channels"),
-                                    attention_type=pixart_config.get("attention_type")
-                                )
-    # controlnet = init_controlnet(controlnet, causal_dit)
-    del transformer
-    transformer_lora_config = LoraConfig(
-            r=lora_rank,
-            lora_alpha=lora_rank,
-            # use_dora=True,
-            init_lora_weights="gaussian",
-            target_modules=["to_k",
-                "to_q",
-                "to_v",
-                "to_out.0",
-                "proj_in",
-                "proj_out",
-                "ff.net.0.proj",
-                "ff.net.2",
-                "proj",
-                "linear",
-                "linear_1",
-                "linear_2"],#ff.net.0.proj ff.net.2
-        )
-    causal_dit.add_adapter(transformer_lora_config)
+lora_state_dict = torch.load(os.path.join(model_global_path, 'shadow_ckpt', 'transformer_lora_pos.bin'), map_location='cpu')
+causal_dit.load_state_dict(lora_state_dict, strict=False)
+controlnet_state_dict = torch.load(os.path.join(model_global_path, 'shadow_ckpt', 'controlnet.bin'), map_location='cpu')
+controlnet.load_state_dict(controlnet_state_dict, strict=True)
 
-    
-    lora_state_dict = torch.load(os.path.join(model_global_path, 'shadow_ckpt', 'transformer_lora_pos.bin'), map_location='cpu')
-    causal_dit.load_state_dict(lora_state_dict, strict=False)
-    controlnet_state_dict = torch.load(os.path.join(model_global_path, 'shadow_ckpt', 'controlnet.bin'), map_location='cpu')
-    controlnet.load_state_dict(controlnet_state_dict, strict=True)
-
-    causal_dit.to('cuda', dtype=weight_dtype)
-    controlnet.to('cuda', dtype=weight_dtype)
-
-    pipeline = CobraPixArtAlphaPipeline.from_pretrained(
-            pretrained_model_name_or_path,
-            transformer=causal_dit,
-            controlnet=controlnet,
-            safety_checker=None,    
-            revision=None,
-            variant=None,
-            torch_dtype=weight_dtype,
-        )
+causal_dit.to('cuda', dtype=weight_dtype)
+controlnet.to('cuda', dtype=weight_dtype)
 
-    pipeline = pipeline.to("cuda")
+pipeline = CobraPixArtAlphaPipeline.from_pretrained(
+        pretrained_model_name_or_path,
+        transformer=causal_dit,
+        controlnet=controlnet,
+        safety_checker=None,    
+        revision=None,
+        variant=None,
+        torch_dtype=weight_dtype,
+    )
 
-    print('loaded pipeline')
+pipeline = pipeline.to("cuda")
 
+print('loaded pipeline')
 
-load_ckpt()
 
 @spaces.GPU
 def change_ckpt(style):