修改计算类型

app.py

@@ -28,10 +28,14 @@ lora_base_path = "./models"
 pipe = FluxPipeline.from_pretrained(base_path, torch_dtype=torch.bfloat16)
 transformer = FluxTransformer2DModel.from_pretrained(base_path, subfolder="transformer", torch_dtype=torch.bfloat16)
 pipe.transformer = transformer
-pipe.to("cuda")
-# 在初始化模型后立即清理GPU缓存和启用注意力切片
-torch.cuda.empty_cache()  # 清理GPU缓存
-pipe.enable_attention_slicing()  # 启用注意力切片以减少内存使用
+try:
+    pipe.to("cuda")
+    # 在初始化模型后立即清理GPU缓存和启用注意力切片
+    torch.cuda.empty_cache()  # 清理GPU缓存
+    pipe.enable_attention_slicing()  # 启用注意力切片以减少内存使用
+except torch.cuda.OutOfMemoryError:
+    print("CUDA out of memory. Switching to CPU.")
+    pipe.to("cpu")
 
 def clear_cache(transformer):
     for name, attn_processor in transformer.attn_processors.items():