Spaces:
Running on Zero
Running on Zero
| import spaces | |
| import gradio as gr | |
| import torch | |
| from diffusers import BriaFiboEditPipeline | |
| from PIL import Image | |
| import tempfile | |
| # Load the model on CPU at module level (ZeroGPU moves to CUDA inside @spaces.GPU) | |
| print("Loading Fibo-Edit-RMBG model...") | |
| pipe = BriaFiboEditPipeline.from_pretrained( | |
| "briaai/Fibo-Edit-RMBG", | |
| torch_dtype=torch.bfloat16, | |
| ) | |
| print("Model loaded on CPU") | |
| INSTRUCTION = {'edit_instruction':'Generate a detailed grayscale alpha matte. Map the opaque foreground to white and the background to black. Produce soft, anti-aliased grayscale gradients at the edges of the subject to represent fine details and transparency.'} | |
| def process(image, num_steps=10, guidance_scale=1.0): | |
| if image is None: | |
| return None, None | |
| pipe.to("cuda") | |
| mask = pipe( | |
| image=image, | |
| prompt=INSTRUCTION, | |
| num_inference_steps=int(num_steps), | |
| guidance_scale=guidance_scale, | |
| ).images[0] | |
| original_rgba = image.convert("RGBA") | |
| mask_gray = mask.convert("L") | |
| if mask_gray.size != original_rgba.size: | |
| mask_gray = mask_gray.resize(original_rgba.size, Image.Resampling.LANCZOS) | |
| original_rgba.putalpha(mask_gray) | |
| result = original_rgba | |
| # Save full-res result for download | |
| tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False) | |
| result.save(tmp.name, format="PNG") | |
| return (image.convert("RGBA"), result), tmp.name | |
| with gr.Blocks(title="Fibo-Edit-RMBG - Background Removal") as demo: | |
| gr.Markdown(""" | |
| # Fibo-Edit-RMBG - Background Removal | |
| Powered by <a href="https://huggingface.co/briaai/Fibo-Edit-RMBG" target="_blank">Bria AI's Fibo-Edit-RMBG model</a> | |
| This model performs background removal on any image | |
| ### How to use: | |
| 1. Upload your image | |
| 2. Adjust settings if needed (optional) | |
| 3. Click Process! | |
| """) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| input_image = gr.Image(label="Input Image", type="pil", height=400) | |
| with gr.Accordion("Advanced Settings", open=False): | |
| num_steps = gr.Slider( | |
| minimum=1, | |
| maximum=50, | |
| value=10, | |
| step=1, | |
| label="Number of Steps", | |
| ) | |
| guidance_scale = gr.Slider( | |
| minimum=1.0, | |
| maximum=15.0, | |
| value=1.0, | |
| step=0.5, | |
| label="Guidance Scale", | |
| ) | |
| process_btn = gr.Button("Remove Background", variant="primary", size="lg") | |
| with gr.Column(scale=1): | |
| slider = gr.ImageSlider( | |
| label="Original vs Background Removed", | |
| type="pil", | |
| image_mode="RGBA", | |
| format="png", | |
| slider_position=50, | |
| max_height=500, | |
| container=True, | |
| ) | |
| download_file = gr.File(label="Download Result (PNG with transparency)") | |
| gr.Markdown(""" | |
| --- | |
| ### About Fibo-Edit-RMBG | |
| Fibo-Edit-RMBG is built on <a href="https://huggingface.co/briaai/Fibo-Edit" target="_blank">Fibo-Edit</a>, an 8B parameter image editing model | |
| that uses structured prompts for precise, deterministic editing. | |
| - **Model**: <a href="https://huggingface.co/briaai/Fibo-Edit-RMBG" target="_blank">briaai/Fibo-Edit-RMBG</a> | |
| - **Architecture**: Based on FIBO with 8B parameters | |
| - **License**: Non-commercial use (<a href="https://creativecommons.org/licenses/by-nc/4.0/" target="_blank">CC BY-NC 4.0</a>) | |
| For commercial use, please <a href="https://bria.ai/contact-us" target="_blank">contact Bria AI</a>. | |
| """) | |
| process_btn.click( | |
| fn=process, | |
| inputs=[input_image, num_steps, guidance_scale], | |
| outputs=[slider, download_file], | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(show_error=True, theme=gr.themes.Soft(primary_hue="sky")) | |