Upload app.py

app.py

@@ -6,6 +6,7 @@ import torch
 from PIL import Image
 from tqdm import tqdm
 import gradio as gr
+import gc  # 导入垃圾回收模块
 
 from safetensors.torch import save_file
 from src.pipeline import FluxPipeline
@@ -16,43 +17,91 @@ from src.lora_helper import set_single_lora, set_multi_lora, unset_lora
 base_path = "black-forest-labs/FLUX.1-dev"    
 lora_base_path = "./models"
 
+# 设置更低的内存使用限制
+torch.backends.cudnn.benchmark = False  # 关闭 cudnn benchmark 以减少内存占用
 
-pipe = FluxPipeline.from_pretrained(base_path, torch_dtype=torch.float32)  # 使用 float32 而不是 bfloat16，以便在 CPU 上运行
-transformer = FluxTransformer2DModel.from_pretrained(base_path, subfolder="transformer", torch_dtype=torch.float32)
+# 使用较低精度和更保守的加载选项
+pipe = FluxPipeline.from_pretrained(
+    base_path, 
+    torch_dtype=torch.float32,
+    low_cpu_mem_usage=True,  # 启用低内存使用模式
+    use_safetensors=True     # 使用 safetensors 以减少内存使用
+)
+transformer = FluxTransformer2DModel.from_pretrained(
+    base_path, 
+    subfolder="transformer", 
+    torch_dtype=torch.float32,
+    low_cpu_mem_usage=True,
+    use_safetensors=True
+)
 pipe.transformer = transformer
-# 不再使用 cuda，在 CPU 环境运行
-# pipe.to("cuda")
 
 def clear_cache(transformer):
     for name, attn_processor in transformer.attn_processors.items():
         attn_processor.bank_kv.clear()
+    # 手动触发垃圾回收
+    gc.collect()
+    torch.cuda.empty_cache() if torch.cuda.is_available() else None
 
 # Define the Gradio interface
-@spaces.CPU()  # 使用 CPU 装饰器替代 GPU 装饰器
+@spaces.CPU()
 def single_condition_generate_image(prompt, spatial_img, height, width, seed, control_type):
-    # Set the control type
-    if control_type == "Ghibli":
-        lora_path = os.path.join(lora_base_path, "Ghibli.safetensors")
-    set_single_lora(pipe.transformer, lora_path, lora_weights=[1], cond_size=512)
-    
-    # Process the image
-    spatial_imgs = [spatial_img] if spatial_img else []
-    
-    # 由于在 CPU 上运行，可能需要降低一些参数来提高性能
-    image = pipe(
-        prompt,
-        height=int(height),
-        width=int(width),
-        guidance_scale=3.5,
-        num_inference_steps=20,  # 减少推理步骤以在 CPU 上更快运行
-        max_sequence_length=512,
-        generator=torch.Generator().manual_seed(seed),  # 移除 "cpu" 参数，因为在 CPU 上默认就是 CPU 生成器
-        subject_images=[],
-        spatial_images=spatial_imgs,
-        cond_size=512,
-    ).images[0]
-    clear_cache(pipe.transformer)
-    return image
+    try:
+        # 限制图像尺寸，减少内存使用
+        max_dimension = 512  # 设置最大尺寸限制
+        if int(height) > max_dimension or int(width) > max_dimension:
+            aspect_ratio = float(width) / float(height)
+            if aspect_ratio > 1:
+                width = max_dimension
+                height = int(max_dimension / aspect_ratio)
+            else:
+                height = max_dimension
+                width = int(max_dimension * aspect_ratio)
+        
+        # Set the control type
+        if control_type == "Ghibli":
+            lora_path = os.path.join(lora_base_path, "Ghibli.safetensors")
+        set_single_lora(pipe.transformer, lora_path, lora_weights=[1], cond_size=512)
+        
+        # 如果有空间图像，确保其尺寸合理
+        if spatial_img:
+            # 调整空间图像尺寸以减少内存使用
+            max_img_size = 1024
+            if max(spatial_img.size) > max_img_size:
+                ratio = max_img_size / max(spatial_img.size)
+                new_size = (int(spatial_img.size[0] * ratio), int(spatial_img.size[1] * ratio))
+                spatial_img = spatial_img.resize(new_size, Image.LANCZOS)
+            spatial_imgs = [spatial_img]
+        else:
+            spatial_imgs = []
+        
+        # 使用更保守的参数
+        image = pipe(
+            prompt,
+            height=int(height),
+            width=int(width),
+            guidance_scale=3.0,  # 略微降低指导比例
+            num_inference_steps=15,  # 进一步减少推理步骤
+            max_sequence_length=384,  # 减少序列长度
+            generator=torch.Generator().manual_seed(seed),
+            subject_images=[],
+            spatial_images=spatial_imgs,
+            cond_size=384,  # 减小条件尺寸
+        ).images[0]
+        
+        # 清理缓存并回收内存
+        clear_cache(pipe.transformer)
+        gc.collect()
+        
+        return image
+    except Exception as e:
+        # 处理错误并清理内存
+        clear_cache(pipe.transformer)
+        gc.collect()
+        print(f"Error during image generation: {str(e)}")
+        # 返回一个错误图像或消息
+        error_img = Image.new('RGB', (400, 200), color=(255, 255, 255))
+        return error_img
 
 # Define the Gradio interface components
 control_types = ["Ghibli"]