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 ---
 class MockSpaces:
     def GPU(self, duration=0):
         def decorator(func):
@@ -16,6 +16,7 @@ class MockSpaces:
         return decorator
 spaces = MockSpaces()
 
+# --- Theme Setup ---
 colors.steel_blue = colors.Color(
     name="steel_blue",
     c50="#EBF3F8",
@@ -83,6 +84,7 @@ 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())
 
@@ -91,39 +93,52 @@ from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
 from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
 from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
 
-dtype = torch.bfloat16
-
-print("Loading Transformer with split config...")
+# --- 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="auto"
+    device_map=transformer_strategy 
 )
 
 print("Loading Pipeline...")
-
 pipe = QwenImageEditPlusPipeline.from_pretrained(
     "Qwen/Qwen-Image-Edit-2509",
     transformer=transformer_model,
     torch_dtype=dtype,
-    device_map="auto" 
+    device_map=device_strategy # <--- 这里必须是 balanced 或 cpu，不能是 auto
 )
 
-
-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")
-pipe.load_lora_weights("dx8152/Qwen-Image-Edit-2509-Light_restoration", weight_name="移除光影.safetensors", adapter_name="light-restoration")
-pipe.load_lora_weights("dx8152/Qwen-Image-Edit-2509-Relight", weight_name="Qwen-Edit-Relight.safetensors", adapter_name="relight")
-pipe.load_lora_weights("dx8152/Qwen-Edit-2509-Multi-Angle-Lighting", weight_name="多角度灯光-251116.safetensors", adapter_name="multi-angle-lighting")
-pipe.load_lora_weights("tlennon-ie/qwen-edit-skin", weight_name="qwen-edit-skin_1.1_000002750.safetensors", adapter_name="edit-skin")
-pipe.load_lora_weights("lovis93/next-scene-qwen-image-lora-2509", weight_name="next-scene_lora-v2-3000.safetensors", adapter_name="next-scene")
-pipe.load_lora_weights("vafipas663/Qwen-Edit-2509-Upscale-LoRA", weight_name="qwen-edit-enhance_64-v3_000001000.safetensors", adapter_name="upscale-image")
-
-try:
-    pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
-except Exception as e:
-    print(f"Warning: FA3 set skipped: {e}")
+# Only load LoRAs and optimization if on GPU to avoid build errors
+if torch.cuda.device_count() > 0:
+    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")
+    pipe.load_lora_weights("dx8152/Qwen-Image-Edit-2509-Light_restoration", weight_name="移除光影.safetensors", adapter_name="light-restoration")
+    pipe.load_lora_weights("dx8152/Qwen-Image-Edit-2509-Relight", weight_name="Qwen-Edit-Relight.safetensors", adapter_name="relight")
+    pipe.load_lora_weights("dx8152/Qwen-Edit-2509-Multi-Angle-Lighting", weight_name="多角度灯光-251116.safetensors", adapter_name="multi-angle-lighting")
+    pipe.load_lora_weights("tlennon-ie/qwen-edit-skin", weight_name="qwen-edit-skin_1.1_000002750.safetensors", adapter_name="edit-skin")
+    pipe.load_lora_weights("lovis93/next-scene-qwen-image-lora-2509", weight_name="next-scene_lora-v2-3000.safetensors", adapter_name="next-scene")
+    pipe.load_lora_weights("vafipas663/Qwen-Edit-2509-Upscale-LoRA", weight_name="qwen-edit-enhance_64-v3_000001000.safetensors", adapter_name="upscale-image")
+
+    try:
+        pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
+    except Exception as e:
+        print(f"Warning: FA3 set skipped: {e}")
 
 MAX_SEED = np.iinfo(np.int32).max
 
@@ -148,6 +163,10 @@ def infer(input_image, prompt, lora_adapter, seed, randomize_seed, guidance_scal
     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",
@@ -185,6 +204,9 @@ 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
     input_pil = input_image.convert("RGB")
     result, seed = infer(input_pil, prompt, lora_adapter, 0, True, 1.0, 4)
     return result, seed
@@ -196,8 +218,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 Fixed)**", elem_id="main-title")
-        gr.Markdown("Forked & Optimized for RunPod 2xGPU Split.")
+        gr.Markdown("# **Qwen-Image-Edit-2509-LoRAs-Fast (2xA40 Ready)**", elem_id="main-title")
         
         with gr.Row(equal_height=True):
             with gr.Column():