|
|
import spaces |
|
|
import torch |
|
|
from diffusers import BriaPipeline |
|
|
from PIL import Image |
|
|
import gradio as gr |
|
|
|
|
|
MODEL_ID = "briaai/BRIA-3.2" |
|
|
pipe = BriaPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16) |
|
|
pipe.to("cuda") |
|
|
|
|
|
@spaces.GPU(duration=1500) |
|
|
def compile_transformer(): |
|
|
with spaces.aoti_capture(pipe.transformer) as call: |
|
|
pipe("arbitrary example prompt") |
|
|
|
|
|
exported = torch.export.export( |
|
|
pipe.transformer, |
|
|
args=call.args, |
|
|
kwargs=call.kwargs, |
|
|
) |
|
|
return spaces.aoti_compile(exported) |
|
|
|
|
|
compiled_transformer = compile_transformer() |
|
|
spaces.aoti_apply(compiled_transformer, pipe.transformer) |
|
|
|
|
|
@spaces.GPU(duration=60) |
|
|
def generate_image(prompt, seed=0): |
|
|
torch.manual_seed(seed) |
|
|
image = pipe(prompt).images[0] |
|
|
return image |
|
|
|
|
|
with gr.Blocks() as demo: |
|
|
gr.Markdown("# BRIA-3.2 Text-to-Image Generator") |
|
|
gr.Markdown("Generate images from text prompts using the BRIA-3.2 model.") |
|
|
|
|
|
with gr.Row(): |
|
|
prompt = gr.Textbox( |
|
|
label="Prompt", |
|
|
value="a cat sitting on a chair", |
|
|
interactive=True |
|
|
) |
|
|
seed = gr.Number( |
|
|
label="Seed (0 for random)", |
|
|
value=0, |
|
|
precision=0 |
|
|
) |
|
|
|
|
|
generate_btn = gr.Button("Generate Image") |
|
|
output = gr.Image(label="Generated Image", type="pil") |
|
|
|
|
|
generate_btn.click( |
|
|
fn=generate_image, |
|
|
inputs=[prompt, seed], |
|
|
outputs=output |
|
|
) |
|
|
|
|
|
gr.Examples( |
|
|
examples=[ |
|
|
["a futuristic cityscape at sunset"], |
|
|
["a forest with glowing mushrooms"], |
|
|
["a steampunk robot drinking tea"], |
|
|
["an astronaut riding a horse on mars"] |
|
|
], |
|
|
inputs=[prompt], |
|
|
outputs=output, |
|
|
fn=generate_image, |
|
|
cache_examples="lazy" |
|
|
) |
|
|
|
|
|
demo.launch() |
