app.py

import gradio as gr
from optimum.intel import OVDiffusionPipeline
import torch

# Load the FLUX.1-schnell model optimized with OpenVINO (INT4)
model_id = "OpenVINO/FLUX.1-schnell-int4-ov"
pipeline = OVDiffusionPipeline.from_pretrained(model_id, device="CPU")

# Define the image generation function
def generate_image(prompt):
    # Generate the image using the pipeline
    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"]],
)

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