Spaces:

---

Duplicated from  zerogpu-aoti/Flux-Compiled-Graph

zerogpu-aoti

/

Flux-Use-Compiled-Graph

like 1

Running on Zero

I don't know how. I just ran the demo with a prompt and it worked.

Sorry I was referring to new compiled model base + lora , will it work with base +lora ?

model  => rahul7star/flux-aot.   

Just used below steps to upload

pipe = DiffusionPipeline.from_pretrained("black-forest-labs/Flux.1-Dev", torch_dtype=dtype).to(device)
pipe.load_lora_weights("rahul7star/ra3hul")
pipe.transformer.set_attn_processor(FlashFluxAttnProcessor3_0())

not sure why shape messing up hmm did the lora didnt make it to GPU lol during swap

in general AOT compilation during runtime works very well , tested same lora

pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev",torch_dtype=torch.bfloat16)
pipe.load_lora_weights("rahul7star/ra3hul")
pipe.to('cuda')

def compile_transformer():
    """Compile the transformer for faster inference"""
    with spaces.aoti_capture(pipe.transformer) as call:
        pipe("optimization test prompt", num_inference_steps=28)
    
    exported = torch.export.export(
        pipe.transformer,
        args=call.args,
        kwargs=call.kwargs,
    )
    return spaces.aoti_compile(exported)

# Apply compiled model during startup
print("Compiling model for optimal performance...")
compiled_transformer = compile_transformer()
spaces.aoti_apply(compiled_transformer, pipe.transformer)
print("✅ Model compiled successfully!")