from wandb.integration.diffusers import autolog
from diffusers import DiffusionPipeline
import torch
from config import PROJECT_NAME
autolog(init=dict(project=PROJECT_NAME))

class SDXLLoraInference:
    """
    Class for running inference using the SDXL-LoRA model to generate stunning product photographs.
    
    Args:
        num_inference_steps (int): The number of inference steps to perform.
        guidance_scale (float): The scale factor for guidance during inference.
    """
    def __init__(self, num_inference_steps: int, guidance_scale: float) -> None:
        self.model_path = "VikramSingh178/sdxl-lora-finetune-product-caption"
        self.pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)
        self.pipe.to("cuda")
        self.pipe.load_lora_weights(self.model_path)
        self.num_inference_steps = num_inference_steps
        self.guidance_scale = guidance_scale

    def run_inference(self, prompt):
        """
        Runs inference using the SDXL-LoRA model to generate a stunning product photograph.
        
        Args:
            prompt: The input prompt for generating the product photograph.
        
        Returns:
            images: The generated product photograph(s).
        """
        
        prompt = prompt
        images = self.pipe(prompt, num_inference_steps=self.num_inference_steps, guidance_scale=self.guidance_scale).images
        return images

inference = SDXLLoraInference(num_inference_steps=100, guidance_scale=2.5)
inference.run_inference(prompt= "A stunning 4k Shot of a Balenciaga X Anime Hoodie with a person wearing it in a party" )