Initial Commit

app.py

@@ -12,6 +12,7 @@ import logging
 import sys
 
 from qwen2_vl.modeling_qwen2_vl import Qwen2VLSimplifiedModel
+from huggingface_hub import snapshot_download
 import spaces
 
 # 设置日志
@@ -25,6 +26,27 @@ logging.basicConfig(
 logger = logging.getLogger(__name__)
 
 MODEL_ID = "Djrango/Qwen2vl-Flux"
+MODEL_CACHE_DIR = "model_cache"
+
+# 预下载所有模型
+def download_models():
+    logger.info("Starting model download...")
+    try:
+        # 下载完整模型仓库
+        snapshot_download(
+            repo_id=MODEL_ID,
+            local_dir=MODEL_CACHE_DIR,
+            local_dir_use_symlinks=False
+        )
+        
+        logger.info("Model download completed successfully")
+    except Exception as e:
+        logger.error(f"Error downloading models: {str(e)}")
+        raise
+
+# 在脚本开始时下载模型
+if not os.path.exists(MODEL_CACHE_DIR):
+    download_models()
 
 # Add aspect ratio options
 ASPECT_RATIOS = {
@@ -61,33 +83,31 @@ class FluxInterface:
         torch.cuda.max_memory_allocated = lambda *args, **kwargs: 0  # 忽略已分配内存的限制
         
         # Load FLUX components
-        tokenizer = CLIPTokenizer.from_pretrained(self.MODEL_ID, subfolder="flux/tokenizer")
-        text_encoder = CLIPTextModel.from_pretrained(self.MODEL_ID, subfolder="flux/text_encoder").to(self.dtype).to(self.device)
-        text_encoder_two = T5EncoderModel.from_pretrained(self.MODEL_ID, subfolder="flux/text_encoder_2").to(self.dtype).to(self.device)
-        tokenizer_two = T5TokenizerFast.from_pretrained(self.MODEL_ID, subfolder="flux/tokenizer_2")
+        tokenizer = CLIPTokenizer.from_pretrained(os.path.join(MODEL_CACHE_DIR, "flux/tokenizer"))
+        text_encoder = CLIPTextModel.from_pretrained(os.path.join(MODEL_CACHE_DIR, "flux/text_encoder")).to(self.dtype).to(self.device)
+        text_encoder_two = T5EncoderModel.from_pretrained(os.path.join(MODEL_CACHE_DIR, "flux/text_encoder_2")).to(self.dtype).to(self.device)
+        tokenizer_two = T5TokenizerFast.from_pretrained(os.path.join(MODEL_CACHE_DIR, "flux/tokenizer_2"))
         
         # Load VAE and transformer
-        vae = AutoencoderKL.from_pretrained(self.MODEL_ID, subfolder="flux/vae").to(self.dtype).to(self.device)
-        transformer = FluxTransformer2DModel.from_pretrained(self.MODEL_ID, subfolder="flux/transformer").to(self.dtype).to(self.device)
-        scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(self.MODEL_ID, subfolder="flux/scheduler", shift=1)
+        vae = AutoencoderKL.from_pretrained(os.path.join(MODEL_CACHE_DIR, "flux/vae")).to(self.dtype).to(self.device)
+        transformer = FluxTransformer2DModel.from_pretrained(os.path.join(MODEL_CACHE_DIR, "flux/transformer")).to(self.dtype).to(self.device)
+        scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(os.path.join(MODEL_CACHE_DIR, "flux/scheduler"), shift=1)
         
         # Load Qwen2VL components
-        qwen2vl = Qwen2VLSimplifiedModel.from_pretrained(self.MODEL_ID, subfolder="qwen2-vl").to(self.dtype).to(self.device)
+        qwen2vl = Qwen2VLSimplifiedModel.from_pretrained(os.path.join(MODEL_CACHE_DIR, "qwen2-vl")).to(self.dtype).to(self.device)
         
-        # Load connector
+        # 加载 connector
         connector = Qwen2Connector().to(self.dtype).to(self.device)
-        connector_path = f"https://huggingface.co/{self.MODEL_ID}/resolve/main/qwen2-vl/connector.pt"
-        connector_state = torch.hub.load_state_dict_from_url(connector_path, map_location='cpu')
-        # Move state dict to dtype before loading
+        connector_path = os.path.join(MODEL_CACHE_DIR, "qwen2-vl/connector.pt")
+        connector_state = torch.load(connector_path, map_location='cpu')
         connector_state = {k: v.to(self.dtype) for k, v in connector_state.items()}
         connector.load_state_dict(connector_state)
         connector = connector.to(self.device)
-        
-        # Load T5 embedder
+    
+        # 加载 T5 embedder
         self.t5_context_embedder = nn.Linear(4096, 3072).to(self.dtype).to(self.device)
-        t5_embedder_path = f"https://huggingface.co/{self.MODEL_ID}/resolve/main/qwen2-vl/t5_embedder.pt"
-        t5_embedder_state = torch.hub.load_state_dict_from_url(t5_embedder_path, map_location='cpu')
-        # Move state dict to dtype before loading
+        t5_embedder_path = os.path.join(MODEL_CACHE_DIR, "qwen2-vl/t5_embedder.pt")
+        t5_embedder_state = torch.load(t5_embedder_path, map_location='cpu')
         t5_embedder_state = {k: v.to(self.dtype) for k, v in t5_embedder_state.items()}
         self.t5_context_embedder.load_state_dict(t5_embedder_state)
         self.t5_context_embedder = self.t5_context_embedder.to(self.device)

requirements.txt

@@ -11,4 +11,5 @@ numpy>=1.24.0
 # Utilities
 protobuf==4.23.4
 sentencepiece==0.2.0
-gradio==5.6.0
+gradio==5.6.0
+huggingface-hub