Create inference.py

inference.py

@@ -0,0 +1,65 @@
+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,
+        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 painting of the 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 = "/home/work/daehyun/Style-Portrait-Generator/scripts/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
+    )