diff --git a/worker/generator.py b/worker/generator.py index 4d42fa6..700d332 100755 --- a/worker/generator.py +++ b/worker/generator.py @@ -60,25 +60,37 @@ class SDXLGenerator: try: dtype = torch.float16 if config.USE_FP16 else torch.float32 - # Use DDIM scheduler for IP-Adapter compatibility - noise_scheduler = DDIMScheduler( - num_train_timesteps=1000, - beta_start=0.00085, - beta_end=0.012, - beta_schedule="scaled_linear", - clip_sample=False, - set_alpha_to_one=False, - steps_offset=1, - ) - - self.pipeline = StableDiffusionXLPipeline.from_pretrained( - config.SDXL_MODEL_ID, - torch_dtype=dtype, - scheduler=noise_scheduler, - use_safetensors=True, - cache_dir=str(config.MODELS_DIR), - add_watermarker=False, - ) + # Only use DDIM scheduler when IP-Adapter is enabled (it requires DDIM) + # Otherwise let the model use its default scheduler for best quality + if config.ENABLE_IP_ADAPTER: + noise_scheduler = DDIMScheduler( + num_train_timesteps=1000, + beta_start=0.00085, + beta_end=0.012, + beta_schedule="scaled_linear", + clip_sample=False, + set_alpha_to_one=False, + steps_offset=1, + ) + logger.info("Using DDIM scheduler for IP-Adapter compatibility") + + self.pipeline = StableDiffusionXLPipeline.from_pretrained( + config.SDXL_MODEL_ID, + torch_dtype=dtype, + scheduler=noise_scheduler, + use_safetensors=True, + cache_dir=str(config.MODELS_DIR), + add_watermarker=False, + ) + else: + # Use model's default scheduler for best quality + self.pipeline = StableDiffusionXLPipeline.from_pretrained( + config.SDXL_MODEL_ID, + torch_dtype=dtype, + use_safetensors=True, + cache_dir=str(config.MODELS_DIR), + add_watermarker=False, + ) # Move to device self.pipeline = self.pipeline.to(self.device)