Update app.py

app.py

@@ -67,18 +67,17 @@ def enable_lora(lora_scale, lora_in, lora_add):
         url = f'https://huggingface.co/{lora_add}/tree/main'
         lora_name = scrape_lora_link(url)    
         pipe.lora_lora_weights(lora_add, weight_name=lora_name)
-        pipe.fuse_lora(lora_scale=lora_scale)
-        pipe.to(device="cuda", dtype=torch.bfloat16)
+        pipe.fuse_lora(lora_scale=lora_scale)   
     
 @spaces.GPU()
 def generate_image(
-    prompt,
-    width=768,
-    height=1024,
-    scales=5,
-    steps=24,
-    seed=-1,
-    nums=1,
+    prompt:str,
+    width:int=768,
+    height:int=1024,
+    scales:float=3.5,
+    steps:int=24,
+    seed:int=-1,
+    nums:int=1,
     progress=gr.Progress(track_tqdm=True)):
 
     if seed == -1:
@@ -107,10 +106,10 @@ def generate_image(
 
 
 def gen(
-    prompt,
+    prompt:str,
     width:int=768,
     height:int=1024,
-    scales:int=5,
+    scales:float=3.5,
     steps:int=24,
     seed:int=-1,
     nums:int=1,
@@ -119,6 +118,7 @@ def gen(
     lora_add:str=""
 ):
     enable_lora(lora_scale, lora_in, lora_add)
+    pipe.to(device="cuda", dtype=torch.bfloat16)
     generate_image(prompt,width,height,scales,steps,seed,nums)