修改参数

app.py

@@ -29,6 +29,9 @@ 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()  # 启用注意力切片以减少内存使用
 
 def clear_cache(transformer):
     for name, attn_processor in transformer.attn_processors.items():
@@ -37,28 +40,27 @@ def clear_cache(transformer):
 # Define the Gradio interface
 @spaces.GPU()
 def single_condition_generate_image(prompt, spatial_img, height, width, seed, control_type):
-    with torch.no_grad():
     # 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 []
-            image = pipe(
-                prompt,
-                height=int(height),
-                width=int(width),
-                guidance_scale=3.5,
-                num_inference_steps=25,
-                max_sequence_length=512,
-                generator=torch.Generator("cpu").manual_seed(seed), 
-                subject_images=[],
-                spatial_images=spatial_imgs,
-                cond_size=512,
-            ).images[0]
-            clear_cache(pipe.transformer)
-            return image
+    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 []
+        image = pipe(
+            prompt,
+            height=int(height),
+            width=int(width),
+            guidance_scale=3.5,
+            num_inference_steps=25,
+            max_sequence_length=512,
+            generator=torch.Generator("cpu").manual_seed(seed), 
+            subject_images=[],
+            spatial_images=spatial_imgs,
+            cond_size=512,
+        ).images[0]
+        clear_cache(pipe.transformer)
+        return image
 
 # Define the Gradio interface components
 control_types = ["Ghibli"]