Spaces:
Build error
Build error
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" ) | |