import gradio as gr
from diffusers import FluxPipeline
import torch

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell",
    torch_dtype=torch.bfloat16
)
pipe.enable_model_cpu_offload()

def generate(prompt):
    image = pipe(prompt, num_inference_steps=4, guidance_scale=0.0).images[0]
    return image

gr.Interface(
    fn=generate,
    inputs="textbox",
    outputs="image",
    title="MustiAI Image Gen",
    allow_flagging="never"
).queue().launch(server_name="0.0.0.0", server_port=7860)