Update optimized.py

optimized.py

@@ -1,10 +1,11 @@
 import torch
 import spaces
 import os
+import diffusers
 from diffusers.utils import load_image
 from diffusers import FluxControlNetModel, FluxControlNetPipeline, AutoencoderKL
 import gradio as gr
-from accelerate import init_empty_weights
+from accelerate import dispatch_model, infer_auto_device_map
 
 def self_attention_slicing(module, slice_size=3):
     """Modified from Diffusers' original for Flux compatibility"""
@@ -59,11 +60,27 @@ pipe = FluxControlNetPipeline.from_pretrained(
 )
 print(f"VRAM used: {torch.cuda.memory_allocated()/1e9:.2f}GB")
 # Proper CPU offloading sequence
-pipe.enable_sequential_cpu_offload(
-    device=torch.device("cuda:0"), 
-    execution_device="cuda"
+# device_map = infer_auto_device_map(pipe)
+# pipe = dispatch_model(pipe, device_map=device_map, main_device="cuda")
+
+device_map = infer_auto_device_map(
+    pipe,
+    max_memory={0:"38GB", "cpu":"64GB"},
+    device_types=["cuda", "cpu"]
+)
+pipe = dispatch_model(
+    pipe, 
+    device_map=device_map, 
+    main_device="cuda"
 )
 
+# For Diffusers v0.20+
+pipe.enable_sequential_cpu_offload() 
+# (No parameters needed)
+
+pipe.unet.to(dtype=torch.bfloat16)
+pipe.controlnet.to(dtype=torch.bfloat16)
+pipe.vae.to(dtype=torch.bfloat16)
 # # 2. Then apply custom VAE slicing
 # if getattr(pipe, "vae", None) is not None:
 #     # Method 1: Use official implementation if available