Upload app.py with huggingface_hub

app.py

@@ -8,7 +8,6 @@ import gradio as gr
 import random
 import numpy as np
 from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL
-from peft import PeftModel
 
 dtype = torch.bfloat16
 device = "cuda"
@@ -19,7 +18,7 @@ good_vae = AutoencoderKL.from_pretrained(
     "black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype
 ).to(device)
 
-# Load base FLUX.1-dev with tiny VAE, then apply CMO LoRA adapter
+# Load base FLUX.1-dev with tiny VAE
 BASE_MODEL = "black-forest-labs/FLUX.1-dev"
 CMO_LORA = "Bruece/FLUX.1-dev-CMO"
 
@@ -27,10 +26,9 @@ pipe = DiffusionPipeline.from_pretrained(
     BASE_MODEL, torch_dtype=dtype, vae=taef1
 ).to(device)
 
-# Load and merge the CMO LoRA adapter
-pipe.transformer = PeftModel.from_pretrained(pipe.transformer, CMO_LORA)
-pipe.transformer = pipe.transformer.merge_and_unload()
-pipe = pipe.to(device)
+# Load the CMO LoRA adapter using diffusers' built-in loader (ZeroGPU-safe)
+pipe.load_lora_weights(CMO_LORA)
+pipe.fuse_lora()
 
 torch.cuda.empty_cache()