Update app.py

app.py

@@ -116,18 +116,28 @@ def run_lora(input_image, prompt, cfg_scale, steps, selected_index, randomize_se
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
 
-    image_generator = generate_image(input_image, prompt_mash, steps, seed, cfg_scale, lora_scale, progress)
+    # image_generator = generate_image(input_image, prompt_mash, steps, seed, cfg_scale, lora_scale, progress)
+
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+    final_image =  pipe(
+        image=input_image,
+        prompt=prompt_mash,
+        num_inference_steps=steps,
+        guidance_scale=cfg_scale,
+        generator=generator,
+        joint_attention_kwargs={"scale": lora_scale},
+    )
     
-    # Consume the generator to get the final image
-    final_image = None
-    step_counter = 0
-    for image in image_generator:
-        step_counter+=1
-        final_image = image
-        progress_bar = f'<div class="progress-container"><div class="progress-bar" style="--current: {step_counter}; --total: {steps};"></div></div>'
-        yield image, seed, gr.update(value=progress_bar, visible=True)
+    # # Consume the generator to get the final image
+    # final_image = None
+    # step_counter = 0
+    # for image in image_generator:
+    #     step_counter+=1
+    #     final_image = image
+    #     progress_bar = f'<div class="progress-container"><div class="progress-bar" style="--current: {step_counter}; --total: {steps};"></div></div>'
+    #     yield image, seed, gr.update(value=progress_bar, visible=True)
         
-    yield final_image, seed, gr.update(value=progress_bar, visible=False)
+    yield final_image, seed, gr.update(visible=True)
 
 css = '''
 #gen_btn{height: 100%}