Update app.py

app.py

@@ -62,12 +62,14 @@ def generate(
 
         if not use_vae:  
             pipe = DiffusionPipeline.from_pretrained(model, torch_dtype=torch.float16)
-            pipe.scheduler = sampler.from_config(pipe.scheduler.config)
+            sampler_class = getattr(spaces, sampler)
+            pipe.scheduler = sampler_class.from_config(pipe.scheduler.config)
             
         if use_vae:
             vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
             pipe = DiffusionPipeline.from_pretrained(model, vae=vae, torch_dtype=torch.float16)
-            pipe.scheduler = sampler.from_config(pipe.scheduler.config)
+            sampler_class = getattr(spaces, sampler)
+            pipe.scheduler = sampler_class.from_config(pipe.scheduler.config)
             
         if use_lora:
             pipe.load_lora_weights(lora)