from optimum.intel import OVStableDiffusionPipeline  # Changed from OVPipelineForText2Image
from openvino import Core, Model  # Updated to use non-deprecated openvino module
import gradio as gr

# Load the model
model_id = "AIFunOver/FLUX.1-dev-openvino-fp16"
try:
    model = OVStableDiffusionPipeline.from_pretrained(
        model_id,
        export=True,  # Enable export to OpenVINO IR if needed
        device="AUTO"  # Use AUTO for flexible device selection
    )
    # Save the model after export to avoid re-conversion
    model.save_pretrained("./flux1_dev_openvino_fp16")
except Exception as e:
    print(f"Error loading model: {e}")
    exit(1)

# Define the image generation function
def generate_image(prompt, num_inference_steps=50, guidance_scale=7.5):
    try:
        # Generate image using the model
        result = model(
            prompt=prompt,
            num_inference_steps=int(num_inference_steps),  # Ensure integer input
            guidance_scale=float(guidance_scale),  # Ensure float input
            height=512,
            width=512
        )
        image = result.images[0]  # Extract the generated image
        return image
    except Exception as e:
        return f"Error generating image: {e}"

# Set up Gradio interface
interface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Enter your image description here..."),
        gr.Slider(label="Inference Steps", minimum=1, maximum=100, value=50, step=1),
        gr.Slider(label="Guidance Scale", minimum=1.0, maximum=20.0, value=7.5, step=0.1)
    ],
    outputs=gr.Image(label="Generated Image"),
    title="FLUX.1-dev Image Generation with OpenVINO",
    description="Enter a text prompt to generate an image using the FLUX.1-dev model optimized with OpenVINO."
)

# Launch the Gradio app
if __name__ == "__main__":
    try:
        interface.launch(server_name="0.0.0.0", server_port=7860)  # Explicit server settings
    except Exception as e:
        print(f"Error launching Gradio interface: {e}")