Initial Commit

app.py

@@ -15,6 +15,7 @@ import os
 # 设置环境变量，强制禁用 accelerate 的显存管理
 os.environ["ACCELERATE_USE_MEMORY_EFFICIENT_ATTENTION"] = "false"
 os.environ["ACCELERATE_DISABLE_MEMORY_EFFICIENT_ATTENTION"] = "1"
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True,garbage_collection_threshold:0.6,max_split_size_mb:512"
 
 from qwen2_vl.modeling_qwen2_vl import Qwen2VLSimplifiedModel
 
@@ -60,6 +61,9 @@ class FluxInterface:
             return
 
         logger.info("Starting model loading...")
+        # 3. 显式设置 PyTorch 缓存分配器的行为
+        torch.cuda.set_per_process_memory_fraction(0.95)  # 允许使用95%的显存
+        torch.cuda.max_memory_allocated = lambda *args, **kwargs: 0  # 忽略已分配内存的限制
         
         # Load FLUX components
         tokenizer = CLIPTokenizer.from_pretrained(self.MODEL_ID, subfolder="flux/tokenizer")