Update app.py

app.py

@@ -34,7 +34,7 @@ pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef
 
 MAX_SEED = 2**32-1
 
-pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
+# pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
 
 class calculateDuration:
     def __init__(self, activity_name=""):
@@ -75,13 +75,13 @@ def upload_image_to_r2(image, account_id, access_key, secret_key, bucket_name):
     return image_file
 
 
-@spaces.GPU(duration=70)
+@spaces.GPU
 def generate_image(prompt, steps, seed, cfg_scale, width, height, lora_scale, progress):
     pipe.to("cuda")
     generator = torch.Generator(device="cuda").manual_seed(seed)
     with calculateDuration("Generating image"):
         # Generate image
-        for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
+        generate_image = pipe(
             prompt=prompt,
             num_inference_steps=steps,
             guidance_scale=cfg_scale,
@@ -91,8 +91,8 @@ def generate_image(prompt, steps, seed, cfg_scale, width, height, lora_scale, pr
             joint_attention_kwargs={"scale": lora_scale},
             output_type="pil",
             good_vae=good_vae,
-        ):
-            yield img
+        )
+    return generate_image
 
 
 def run_lora(prompt, cfg_scale, steps, lora_repo, lora_name, randomize_seed, seed, width, height, lora_scale, upload_to_r2, account_id, access_key, secret_key, bucket, progress=gr.Progress(track_tqdm=True)):
@@ -106,21 +106,10 @@ def run_lora(prompt, cfg_scale, steps, lora_repo, lora_name, randomize_seed, see
             pipe.load_lora_weights(lora_repo, weight_name=lora_name)
 
     # Set random seed for reproducibility
-    with calculateDuration("Randomizing seed"):
-        if randomize_seed:
-            seed = random.randint(0, MAX_SEED)
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
 
-    image_generator = generate_image(prompt, steps, seed, cfg_scale, width, height, lora_scale, progress)
-    
-    # Consume the generator to get the final image
-    final_image = None
-    step_counter = 0
-    final_image = None
-    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),  json.dumps({"status": "processing"})
+    final_image = generate_image(prompt, steps, seed, cfg_scale, width, height, lora_scale, progress)
     
     if upload_to_r2:
         url = upload_image_to_r2(final_image, account_id, access_key, secret_key, bucket)