Update app.py

app.py

@@ -106,19 +106,21 @@ base_model_path: str = "stabilityai/stable-diffusion-xl-base-1.0"
 refiner_model_path: str = "stabilityai/stable-diffusion-xl-refiner-1.0" 
 unet_path: str = "cosmicman/CosmicMan-SDXL"
 SCHEDULER = schedule_map[schedule]
-scheduler = SCHEDULER.from_pretrained(base_model_path, subfolder="scheduler")
-# unet = UNet2DConditionModel.from_pretrained(unet_path)
+scheduler = SCHEDULER.from_pretrained(base_model_path, subfolder="scheduler", torch_dtype=torch.float16)
+# unet = UNet2DConditionModel.from_pretrained(unet_path, torch_dtype=torch.float16)
 
 pipe = StableDiffusionXLPipeline.from_pretrained(
     base_model_path,
     # unet=unet,
     scheduler=scheduler,
+    torch_dtype=torch.float16,
     use_safetensors=True
 ).to("cuda")
 pipe.watermark = NoWatermark()
 refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained(
     base_model_path, # we found use base_model_path instead of refiner_model_path may get a better performance
     scheduler=scheduler,
+    ,torch_dtype=torch.float16,
     use_safetensors=True
 ).to("cuda")
 refiner.watermark = NoWatermark()
@@ -184,7 +186,7 @@ with gr.Blocks(theme=gr.themes.Soft(),css="style.css") as demo:
                 value=0,
             )
             random_seed = gr.Checkbox(label="Randomize seed", value=True)
-            img_num = gr.Slider(minimum=1, maximum=4, value=4, label="Number of images", step=1)   
+            img_num = gr.Slider(minimum=1, maximum=4, value=1, label="Number of images", step=1)   
 
     gr.Examples(
         examples=examples,