app.py

import gradio as gr
from optimum.intel import OVDiffusionPipeline
from threading import Lock

# Load the pipeline globally (includes the correct tokenizer)
model_id = "OpenVINO/FLUX.1-schnell-int4-ov"
pipeline = OVDiffusionPipeline.from_pretrained(model_id, device="CPU")
lock = Lock()

# Define the image generation function
def generate_image(prompt):
    with lock:
        image = pipeline(prompt, num_inference_steps=4, guidance_scale=3.5).images[0]
    return image

# Create the Gradio interface
interface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Enter your prompt", placeholder="e.g., A futuristic cityscape at sunset"),
    outputs=gr.Image(label="Generated Image"),
    title="FLUX.1-Schnell (OpenVINO INT4) Image Generator",
    description="Generate images from text prompts using FLUX.1-schnell optimized for CPU with OpenVINO.",
    examples=[["A serene mountain landscape"], ["A cyberpunk city at night"]],
    cache_examples=False
)

# Launch the interface
if __name__ == "__main__":
    interface.launch()