Update app.py

app.py

@@ -1,8 +1,9 @@
 import torch
 import gradio as gr
-from diffusers import DiffusionPipeline, AutoPipelineForImage2Image, AutoencoderTiny, AutoencoderKL
+from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL
 from diffusers.utils import load_image
 import os
+import gc
 from PIL import Image
 import time
 
@@ -21,23 +22,20 @@ lora_usage = {lora["title"]: 0 for lora in loras}
 device = "cpu"
 dtype = torch.float32  # Use float32 for CPU compatibility
 
-# Initialize pipelines
+# Initialize a single pipeline with CPU offloading
 base_model = "black-forest-labs/FLUX.1-dev"
 taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype)
 good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype)
 
-pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1)
-pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
+pipe = DiffusionPipeline.from_pretrained(
     base_model,
-    vae=good_vae,
-    transformer=pipe.transformer,
-    text_encoder=pipe.text_encoder,
-    tokenizer=pipe.tokenizer,
-    text_encoder_2=pipe.text_encoder_2,
-    tokenizer_2=pipe.tokenizer_2,
-    torch_dtype=dtype
+    torch_dtype=dtype,
+    vae=taef1,
 )
 
+# Enable CPU offloading to reduce memory usage
+pipe.enable_model_cpu_offload()
+
 # Custom CSS
 css = """
 #title {
@@ -84,39 +82,50 @@ def update_lora_info(selected_index, custom_lora):
 def remove_custom_lora(selected_index):
     return None, gr.HTML(visible=False), gr.Button(visible=False), gr.Markdown(value=update_lora_info(selected_index, None)[0])
 
-# Image generation functions
-def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale):
-    generator = torch.Generator(device=device).manual_seed(seed)
-    with calculateDuration("Generating image"):
-        for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
-            prompt=prompt_mash,
-            num_inference_steps=steps,
-            guidance_scale=cfg_scale,
-            width=width,
-            height=height,
-            generator=generator,
-            joint_attention_kwargs={"scale": lora_scale},
-            output_type="pil",
-            good_vae=good_vae,
-        ):
-            yield img
-
-def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps, cfg_scale, width, height, lora_scale, seed):
+# Image generation function (combined for both text-to-image and image-to-image)
+def generate_image(
+    prompt_mash,
+    image_input_path,
+    image_strength,
+    steps,
+    seed,
+    cfg_scale,
+    width,
+    height,
+    lora_scale
+):
     generator = torch.Generator(device=device).manual_seed(seed)
-    image_input = load_image(image_input_path)
-    final_image = pipe_i2i(
-        prompt=prompt_mash,
-        image=image_input,
-        strength=image_strength,
-        num_inference_steps=steps,
-        guidance_scale=cfg_scale,
-        width=width,
-        height=height,
-        generator=generator,
-        joint_attention_kwargs={"scale": lora_scale},
-        output_type="pil",
-    ).images[0]
-    return final_image
+    
+    # Configure pipeline for text-to-image or image-to-image
+    kwargs = {
+        "prompt": prompt_mash,
+        "num_inference_steps": steps,
+        "guidance_scale": cfg_scale,
+        "width": width,
+        "height": height,
+        "generator": generator,
+        "joint_attention_kwargs": {"scale": lora_scale},
+        "output_type": "pil",
+        "good_vae": good_vae,
+    }
+    
+    if image_input_path:
+        image_input = load_image(image_input_path)
+        kwargs.update({
+            "image": image_input,
+            "strength": image_strength,
+        })
+        with calculateDuration("Generating image-to-image"):
+            result = pipe(**kwargs).images[0]
+    else:
+        with calculateDuration("Generating text-to-image"):
+            result = pipe(**kwargs).images[0]
+    
+    # Clear memory after generation
+    torch.cuda.empty_cache()  # No effect on CPU, but harmless
+    gc.collect()
+    
+    return result
 
 def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale):
     global lora_usage
@@ -129,19 +138,27 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
     lora_usage[selected_lora["title"]] += 1
     pipe.unload_lora_weights()
     pipe.load_lora_weights(lora_path)
-    pipe_i2i.load_lora_weights(lora_path)
     if prompt == "":
         prompt = trigger_word
     else:
         prompt_mash = f"{prompt}, {trigger_word}"
     if randomize_seed:
         seed = int(time.time())
-    if image_input:
-        final_image = generate_image_to_image(prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, lora_scale, seed)
-        return final_image, seed, gr.Markdown(value=f"**Seed**: {seed}", visible=True)
-    else:
-        for img in generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale):
-            yield img, seed, gr.Markdown(value=f"**Seed**: {seed}", visible=True)
+    
+    # Generate the image
+    final_image = generate_image(
+        prompt_mash,
+        image_input,
+        image_strength,
+        steps,
+        seed,
+        cfg_scale,
+        width,
+        height,
+        lora_scale
+    )
+    
+    return final_image, seed, gr.Markdown(value=f"**Seed**: {seed}", visible=True)
 
 def generate_usage_chart():
     sorted_usage = sorted(lora_usage.items(), key=lambda x: x[1], reverse=True)[:5]
@@ -239,11 +256,11 @@ with gr.Blocks(theme="YTheme/Minecraft", css=css, delete_cache=(60, 60)) as app:
                 refresh_chart_button = gr.Button("Refresh Usage Chart")
     with gr.Accordion("Advanced Settings", open=False):
         with gr.Row():
-            steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=25, step=1)
+            steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=10, step=1)  # Reduced default steps
             cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, value=7, step=0.1)
         with gr.Row():
-            width = gr.Slider(label="Width", minimum=256, maximum=1024, value=512, step=64)
-            height = gr.Slider(label="Height", minimum=256, maximum=1024, value=512, step=64)
+            width = gr.Slider(label="Width", minimum=256, maximum=1024, value=256, step=64)  # Reduced default resolution
+            height = gr.Slider(label="Height", minimum=256, maximum=1024, value=256, step=64)  # Reduced default resolution
         with gr.Row():
             lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=1, value=0.8, step=0.1)
             image_strength = gr.Slider(label="Image Strength", minimum=0, maximum=1, value=0.5, step=0.1, visible=False)