import torch
from diffusers import FluxPipeline

# Specify the path to your merged model
model_path = "output_checkpoint.safetensors"  # Replace with the actual path

# Load the merged model
pipeline = FluxPipeline.from_pretrained(model_path, torch_dtype=torch.float16)  # Use float32 if you don't have a GPU

# Set the model to evaluation mode
pipeline.eval()

# Example: Generating an image with a prompt
prompt = "A serene landscape with mountains and a lake"  # Customize your prompt here
image = pipeline(prompt).images[0]

# Save or display the generated image
image.save("generated_image.png")