Update app.py

app.py

@@ -5,24 +5,19 @@ import logging
 import torch
 from PIL import Image
 import spaces
-from diffusers import DiffusionPipeline, AutoencoderTiny
-from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
+from diffusers import DiffusionPipeline
 from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, snapshot_download
 import copy
 import random
 import time
 
-dtype = torch.bfloat16
-device = "cuda"
-
 # Load LoRAs from JSON file
 with open('loras.json', 'r') as f:
     loras = json.load(f)
 
 # Initialize the base model
-taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
-pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype, vae=taef1).to(device)
-torch.cuda.empty_cache()
+base_model = "black-forest-labs/FLUX.1-dev"
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
 
 MAX_SEED = 2**32-1
 
@@ -65,22 +60,20 @@ def update_selection(evt: gr.SelectData, width, height):
 @spaces.GPU(duration=70)
 def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, progress):
     pipe.to("cuda")
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device="cuda").manual_seed(seed)
     
     with calculateDuration("Generating image"):
-        for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
+        # Generate image
+        image = pipe(
             prompt=prompt_mash,
             num_inference_steps=steps,
+            guidance_scale=cfg_scale,
             width=width,
             height=height,
-            guidance_scale=cfg_scale,
-            #generator=generator,
-            output_type="pil",
+            generator=generator,
             joint_attention_kwargs={"scale": lora_scale},
-        ):
-            yield img, seed
+        ).images[0]
+    return image
 
 def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
     if selected_index is None: