修改内存大小占用

app.py

@@ -25,8 +25,8 @@ lora_base_path = "./models"
 
 
 # pipe = FluxPipeline.from_pretrained(base_path, torch_dtype=torch.bfloat16)
-pipe = FluxPipeline.from_pretrained(base_path, torch_dtype=torch.bfloat8)
-transformer = FluxTransformer2DModel.from_pretrained(base_path, subfolder="transformer", torch_dtype=torch.bfloat8)
+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")
 
@@ -37,27 +37,28 @@ 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"]