Update app.py

app.py

@@ -14,7 +14,7 @@ from decoder import SketchDecoder
 from transformers import AutoTokenizer, AutoProcessor
 from qwen_vl_utils import process_vision_info
 from tokenizer import SVGTokenizer
-import spaces  # 引入 GPU 装饰器
+import spaces
 
 # 读取配置
 with open('config.yaml', 'r') as f:
@@ -53,19 +53,38 @@ def load_models():
         tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-VL-3B-Instruct", padding_side="left")
         processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-3B-Instruct", padding_side="left")
         
-        sketch_decoder = SketchDecoder()
+        # [关键修改] 1. 下载权重
+        print("Downloading weights...")
         sketch_weight_path = hf_hub_download(repo_id="OmniSVG/OmniSVG", filename="pytorch_model.bin")
         
-        # 加载权重
-        sketch_decoder.load_state_dict(torch.load(sketch_weight_path, map_location="cpu"), strict=False)
+        # [关键修改] 2. 加载权重到 CPU 内存
+        print("Loading state dict...")
+        full_state_dict = torch.load(sketch_weight_path, map_location="cpu")
+        
+        # [关键修改] 3. 处理权重键名 (Strip 'transformer.' prefix)
+        # 原始权重是 SketchDecoder 保存的，带有 "transformer." 前缀
+        # 我们要直接传给内部的 Qwen 模型，所以需要去掉这个前缀
+        qwen_state_dict = {}
+        for key in list(full_state_dict.keys()):
+            if key.startswith("transformer."):
+                new_key = key.replace("transformer.", "", 1)
+                qwen_state_dict[new_key] = full_state_dict.pop(key) # pop saving memory
+        
+        del full_state_dict # 释放内存
+        gc.collect()
+
+        # [关键修改] 4. 初始化 Decoder 并传入处理后的权重
+        # 此时 bitsandbytes 会在初始化时直接将这些权重量化为 4-bit
+        print("Initializing quantized model with custom weights...")
+        sketch_decoder = SketchDecoder(state_dict=qwen_state_dict)
+        
+        # 此时模型已经在 GPU 上（由 device_map="auto" 处理）
+        # 且权重已经是 OmniSVG 的权重了，无需再次 load_state_dict
         
-        if device.type == "cuda":
-            sketch_decoder = sketch_decoder.to(device)
-            
         sketch_decoder.eval()
         svg_tokenizer = SVGTokenizer('config.yaml')
         
-    print(f"✅ Models loaded on {device}")
+    print(f"✅ Models loaded successfully!")
 
 def process_and_resize_image(image_input, target_size=(200, 200)):
     if isinstance(image_input, str):
@@ -264,7 +283,6 @@ def create_interface():
     ]
     example_images = get_example_images()
     
-    # 删除了 theme 参数
     with gr.Blocks(title="OmniSVG Demo Page") as demo:
         gr.Markdown("# OmniSVG Demo Page")
         gr.Markdown("Generate SVG code from images or text descriptions")
@@ -279,7 +297,6 @@ def create_interface():
                         image_generate_btn = gr.Button("Generate SVG", variant="primary")
                     
                     with gr.Column():
-                        # 删除了 show_copy_button=True
                         image_svg_output = gr.Textbox(label="Generated SVG Code", lines=10, max_lines=20)
                         image_png_preview = gr.Image(label="SVG Preview", type="pil")
                 
@@ -293,7 +310,6 @@ def create_interface():
                         text_generate_btn = gr.Button("Generate SVG", variant="primary")
                     
                     with gr.Column():
-                        # 删除了 show_copy_button=True
                         text_svg_output = gr.Textbox(label="Generated SVG Code", lines=10, max_lines=20)
                         text_png_preview = gr.Image(label="SVG Preview", type="pil")