Create test.py

test.py

@@ -0,0 +1,18 @@
+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")