File size: 2,305 Bytes
f69a255 888eca4 f69a255 10ab55c f69a255 10ab55c f69a255 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
import torch
from diffusers import StableDiffusionPipeline
import os
def inference(
model_path: str,
prompt: str,
output_image_path: str = "generated_image.png",
guidance_scale: float = 7.5,
num_inference_steps: int = 50,
use_fp16: bool = True,
):
"""
Generates an image using a fine-tuned Stable Diffusion model.
Args:
model_path (str): The path (or repo ID) to the fine-tuned Stable Diffusion model.
prompt (str): The text prompt used to generate an image.
output_image_path (str): The path to save the generated image.
guidance_scale (float): Classifier-Free Guidance scale.
num_inference_steps (int): Number of inference steps in the diffusion process.
use_fp16 (bool): Whether to use fp16 (half precision). Recommended if a GPU is available.
"""
# Load the model (use half precision if GPU is available)
# If you only have CPU, consider using torch.float32 or omitting torch_dtype
pipe = StableDiffusionPipeline.from_pretrained(
model_path,
subfolder="stable_diffusion-v1-4-wikiart",
torch_dtype=torch.float16 if use_fp16 and torch.cuda.is_available() else torch.float32
)
# Set the device (use GPU if available)
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe.to(device)
# Generate the image
with torch.autocast(device) if (use_fp16 and device == "cuda") else torch.no_grad():
result = pipe(
prompt=prompt,
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps
)
# Save the result
image = result.images[0]
image.save(output_image_path)
print(f"Saved Image: {os.path.abspath(output_image_path)}")
if __name__ == "__main__":
# Example prompt
sample_prompt = "A paint of Eiffel tower, in the style of eric fischl."
# Assuming the fine-tuned model was saved with pipe.save_pretrained("stable-diffusion-wikiart-final")
finetuned_model_path = "stable-diffusion-wikiart-final"
# Run inference
inference(
model_path=finetuned_model_path,
prompt=sample_prompt,
output_image_path="wikiart_inference_result.png",
guidance_scale=7.5,
num_inference_steps=50,
use_fp16=True
) |