Update app.py

app.py

@@ -49,7 +49,10 @@ def update_selection(evt: gr.SelectData):
     )
 
 @spaces.GPU(duration=80)
-def generate_image(prompt, trigger_word, steps, cfg_scale, width, height, generator, lora_scale, progress):
+def generate_image(prompt, trigger_word, steps, seed, cfg_scale, width, height, generator, lora_scale, progress):
+    
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+    
     with calculateDuration("Generating image"):
         # Generate image
         image = pipe(
@@ -82,9 +85,8 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, wid
     with calculateDuration("Randomizing seed"):
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
-        generator = torch.Generator(device="cpu").manual_seed(seed)
     
-    image = generate_image(prompt, trigger_word, steps, cfg_scale, width, height, generator, lora_scale, progress)
+    image = generate_image(prompt, trigger_word, steps, seed, cfg_scale, width, height, generator, lora_scale, progress)
 
     pipe.unload_lora_weights()
     return image, seed