Create inference.py

inference.py

@@ -0,0 +1,64 @@
+import torch
+from diffusers import StableDiffusionXLPipeline, DiffusionPipeline, AutoencoderKL
+from PIL import Image
+from io import BytesIO
+
+class ImageGenerator:
+    def __init__(self):
+        self.model_base = "femboysLover/blue_pencil-fp16-XL"
+        self.v_autoencoder = "madebyollin/sdxl-vae-fp16-fix"
+        self.model_refiner = "stabilityai/stable-diffusion-xl-refiner-1.0"
+
+        # Load the VAE model
+        self.vae = AutoencoderKL.from_pretrained(self.v_autoencoder, torch_dtype=torch.float16)
+
+        # Load the main pipeline
+        self.pipe = StableDiffusionXLPipeline.from_pretrained(
+            self.model_base,
+            torch_dtype=torch.float16,
+            vae=self.vae,
+            add_watermarker=False,
+        )
+        self.pipe.safety_checker = None
+        self.pipe.to("cuda")
+
+        # Load the refiner pipeline
+        self.pipe_refiner = DiffusionPipeline.from_pretrained(
+            self.model_refiner, 
+            torch_dtype=torch.float16, 
+            use_safetensors=True, 
+            variant="fp16"
+        )
+        self.pipe_refiner.enable_model_cpu_offload()
+
+    def generate_image(self, prompt, prompt2, negative_prompt, negative_prompt2, strength=0.3, denoising_start=0.8):
+        # Generate base latent image
+        image_base_latent = self.pipe(prompt).images[0]
+
+        # Refine the image
+        image_refiner = self.pipe_refiner(
+            prompt=prompt,
+            prompt_2=prompt2,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt2,
+            image=image_base_latent,
+            num_inference_steps=25,
+            height=1024,
+            width=1024,
+            strength=strength,
+            denoising_start=denoising_start
+        ).images[0]
+
+        # Convert the image to a format that can be easily outputted (e.g., bytes)
+        buffer = BytesIO()
+        image_refiner.save(buffer, format="JPEG")
+        return buffer.getvalue()
+
+# Usage example
+image_generator = ImageGenerator()
+result = image_generator.generate_image(
+    prompt="A description of the image you want to generate",
+    prompt2="Additional description if needed",
+    negative_prompt="What you want to avoid in the image",
+    negative_prompt2="Additional negative prompt if needed"
+)