Update app.py

app.py

@@ -8,7 +8,7 @@ from typing import Iterable
 from gradio.themes import Soft
 from gradio.themes.utils import colors, fonts, sizes
 
-# --- Mock Spaces ---
+# --- Mock Spaces (保持不变) ---
 class MockSpaces:
     def GPU(self, duration=0):
         def decorator(func):
@@ -16,7 +16,7 @@ class MockSpaces:
         return decorator
 spaces = MockSpaces()
 
-# --- Theme Setup ---
+# --- Theme Setup (保持不变) ---
 colors.steel_blue = colors.Color(
     name="steel_blue",
     c50="#EBF3F8",
@@ -84,47 +84,39 @@ class SteelBlueTheme(Soft):
         )
 steel_blue_theme = SteelBlueTheme()
 
-# --- Debug Info ---
-print("CUDA_VISIBLE_DEVICES:", os.environ.get("CUDA_VISIBLE_DEVICES"))
-print("GPU Count:", torch.cuda.device_count())
-
+# --- 关键修改：按需加载 ---
 from diffusers import FlowMatchEulerDiscreteScheduler
 from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
 from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
 from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
 
-# --- FIX: Determine Device Strategy ---
-# 如果在 HF 构建环境(无 GPU)，使用 CPU 防止报错
-# 如果在 RunPod (有 GPU)，使用 "balanced" 策略 (Pipeline 不支持 "auto")
-if torch.cuda.device_count() > 0:
-    device_strategy = "balanced" # Pipeline level strategy
-    transformer_strategy = "auto" # Transformer level strategy
-    dtype = torch.bfloat16
-    print(f"Running on GPU with strategy: {device_strategy}")
-else:
-    device_strategy = "cpu"
-    transformer_strategy = "cpu"
-    dtype = torch.float32 # CPU usually prefers float32
-    print("Running on CPU (Build Environment detected)")
-
-print("Loading Transformer...")
-transformer_model = QwenImageTransformer2DModel.from_pretrained(
-    "linoyts/Qwen-Image-Edit-Rapid-AIO", 
-    subfolder='transformer',
-    torch_dtype=dtype,
-    device_map=transformer_strategy 
-)
-
-print("Loading Pipeline...")
-pipe = QwenImageEditPlusPipeline.from_pretrained(
-    "Qwen/Qwen-Image-Edit-2509",
-    transformer=transformer_model,
-    torch_dtype=dtype,
-    device_map=device_strategy # <--- 这里必须是 balanced 或 cpu，不能是 auto
-)
+pipe = None # 全局变量初始化为空
 
-# Only load LoRAs and optimization if on GPU to avoid build errors
-if torch.cuda.device_count() > 0:
+# 检测逻辑：只有在真正有 GPU 的时候才加载模型
+# 这样 HF 的 CPU 构建服务器会直接跳过这里，瞬间完成构建
+if torch.cuda.is_available():
+    print("GPU detected! Initializing model for 2x A40 Environment...")
+    dtype = torch.bfloat16
+    
+    # 1. Load Transformer (device_map="auto" for multi-gpu split)
+    print("Loading Transformer...")
+    transformer_model = QwenImageTransformer2DModel.from_pretrained(
+        "linoyts/Qwen-Image-Edit-Rapid-AIO", 
+        subfolder='transformer',
+        torch_dtype=dtype,
+        device_map="auto" 
+    )
+
+    # 2. Load Pipeline (device_map="balanced" compatible with diffusers)
+    print("Loading Pipeline...")
+    pipe = QwenImageEditPlusPipeline.from_pretrained(
+        "Qwen/Qwen-Image-Edit-2509",
+        transformer=transformer_model,
+        torch_dtype=dtype,
+        device_map="balanced"
+    )
+
+    # 3. Load LoRAs
     print("Loading LoRAs...")
     pipe.load_lora_weights("autoweeb/Qwen-Image-Edit-2509-Photo-to-Anime", weight_name="Qwen-Image-Edit-2509-Photo-to-Anime_000001000.safetensors", adapter_name="anime")
     pipe.load_lora_weights("dx8152/Qwen-Edit-2509-Multiple-angles", weight_name="镜头转换.safetensors", adapter_name="multiple-angles")
@@ -139,6 +131,8 @@ if torch.cuda.device_count() > 0:
         pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
     except Exception as e:
         print(f"Warning: FA3 set skipped: {e}")
+else:
+    print("No GPU detected (likely HF Build Environment). SKIPPING MODEL LOAD to save memory.")
 
 MAX_SEED = np.iinfo(np.int32).max
 
@@ -160,13 +154,13 @@ def update_dimensions_on_upload(image):
 
 @spaces.GPU(duration=30)
 def infer(input_image, prompt, lora_adapter, seed, randomize_seed, guidance_scale, steps, progress=gr.Progress(track_tqdm=True)):
+    # 运行时检查：如果 pipe 没加载（说明没 GPU），直接报错
+    if pipe is None:
+        raise gr.Error("Error: Model not loaded. Is a GPU available?")
+
     if input_image is None:
         raise gr.Error("Please upload an image to edit.")
     
-    # 如果没有 GPU (比如在 HF 预览界面)，直接报错提示
-    if torch.cuda.device_count() == 0:
-        raise gr.Error("Running on CPU-only environment. Please run on GPU.")
-
     adapters_map = {
         "Photo-to-Anime": "anime",
         "Multiple-Angles": "multiple-angles",
@@ -204,9 +198,7 @@ def infer(input_image, prompt, lora_adapter, seed, randomize_seed, guidance_scal
 
 @spaces.GPU(duration=30)
 def infer_example(input_image, prompt, lora_adapter):
-    # 如果是 HF 预览构建过程，跳过
-    if torch.cuda.device_count() == 0: 
-        return None, 0
+    if pipe is None: return None, 0
     input_pil = input_image.convert("RGB")
     result, seed = infer(input_pil, prompt, lora_adapter, 0, True, 1.0, 4)
     return result, seed
@@ -218,7 +210,7 @@ css="""
 
 with gr.Blocks(css=css, theme=steel_blue_theme) as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown("# **Qwen-Image-Edit-2509-LoRAs-Fast (2xA40 Ready)**", elem_id="main-title")
+        gr.Markdown("# **Qwen-Image-Edit-2509-LoRAs-Fast (RunPod Optimized)**", elem_id="main-title")
         
         with gr.Row(equal_height=True):
             with gr.Column():