Update app.py

app.py

@@ -5,10 +5,8 @@ import logging
 import torch
 from PIL import Image
 import spaces
-from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL, AutoPipelineForImage2Image
-from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
-from diffusers.utils import load_image
-from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, snapshot_download
+from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL
+from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard
 from transformers import AutoModelForCausalLM, CLIPTokenizer, CLIPProcessor, CLIPModel, LongformerTokenizer, LongformerModel
 import copy
 import random
@@ -47,9 +45,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)
-
-def process_input(input_text):
+ef process_input(input_text):
     # Tokenize and truncate input
     inputs = clip_processor(text=input_text, return_tensors="pt", padding=True, truncation=True, max_length=77)
     return inputs
@@ -93,7 +89,7 @@ def download_file(url, directory=None):
         file.write(response.content)
     
     return filepath
-            
+
 def update_selection(evt: gr.SelectData, selected_indices, loras_state, width, height):
     selected_index = evt.index
     selected_indices = selected_indices or []
@@ -288,39 +284,11 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
     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(
-            prompt=prompt_mash,
-            num_inference_steps=steps,
-            guidance_scale=cfg_scale,
-            width=width,
-            height=height,
-            generator=generator,
-            joint_attention_kwargs={"scale": 1.0},
-            output_type="pil",
-            good_vae=good_vae,
-        ):
+        for img in pipe(prompt=prompt_mash, num_inference_steps=steps, guidance_scale=cfg_scale, width=width, height=height, generator=generator):
             yield img
 
-def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps, cfg_scale, width, height, seed):
-    pipe_i2i.to("cuda")
-    generator = torch.Generator(device="cuda").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": 1.0},
-        output_type="pil",
-    ).images[0]
-    return final_image
-
 @spaces.GPU(duration=75)
-def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, randomize_seed, seed, width, height, loras_state, progress=gr.Progress(track_tqdm=True)):
+def run_lora(prompt, selected_indices, loras_state):
     if not selected_indices:
         raise gr.Error("You must select at least one LoRA before proceeding.")
 
@@ -338,12 +306,7 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
                 appends.append(trigger_word)
     prompt_mash = " ".join(prepends + [prompt] + appends)
     print("Prompt Mash: ", prompt_mash)
-    # Unload previous LoRA weights
-    with calculateDuration("Unloading LoRA"):
-        pipe.unload_lora_weights()
-        pipe_i2i.unload_lora_weights()
-        
-    print(pipe.get_active_adapters())
+    
     # Load LoRA weights with respective scales
     lora_names = []
     lora_weights = []
@@ -352,46 +315,27 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
             lora_name = f"lora_{idx}"
             lora_names.append(lora_name)
             print(f"Lora Name: {lora_name}")
-            lora_weights.append(lora_scale_1 if idx == 0 else lora_scale_2)
+            lora_weights.append(1.15)  # Assuming a default scale
             lora_path = lora['repo']
             weight_name = lora.get("weights")
             print(f"Lora Path: {lora_path}")
-            pipe_to_use = pipe_i2i if image_input is not None else pipe
-            pipe_to_use.load_lora_weights(
+            pipe.load_lora_weights(
                 lora_path, 
                 weight_name=weight_name if weight_name else None,
                 low_cpu_mem_usage=True,
                 adapter_name=lora_name
             )
-#            if image_input is not None: pipe_i2i = pipe_to_use
-#            else: pipe = pipe_to_use
         print("Loaded LoRAs:", lora_names)
         print("Adapter weights:", lora_weights)
-        if image_input is not None:
-            pipe_i2i.set_adapters(lora_names, adapter_weights=lora_weights)
-        else:
-            pipe.set_adapters(lora_names, adapter_weights=lora_weights)
-    print(pipe.get_active_adapters())
+        pipe.set_adapters(lora_names, adapter_weights=lora_weights)
+        print(pipe.get_active_adapters())
+    
     # Set random seed for reproducibility
-    with calculateDuration("Randomizing seed"):
-        if randomize_seed:
-            seed = random.randint(0, MAX_SEED)
+    seed = random.randint(0, MAX_SEED)
 
     # Generate image
-    if image_input is not None:
-        final_image = generate_image_to_image(prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, seed)
-        yield final_image, seed, gr.update(visible=False)
-    else:
-        image_generator = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
-        # 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)
+    final_image = generate_image(prompt_mash, 50, seed, 7.5, 512, 512, None)  # Example parameters
+    yield final_image, seed, gr.update(visible=False)
 
 run_lora.zerogpu = True
 
@@ -451,7 +395,7 @@ def update_history(new_image, history):
     history.insert(0, new_image)
     return history
 
-css = '''
+ccss = '''
 #gen_btn{height: 100%}
 #title{text-align: center}
 #title h1{font-size: 3em; display:inline-flex; align-items:center}
@@ -500,7 +444,7 @@ with gr.Blocks(css=css, delete_cache=(60, 60)) as app:
                 with gr.Column(scale=3, min_width=100):
                     selected_info_1 = gr.Markdown("Select a LoRA 1")
                 with gr.Column(scale=5, min_width=50):
-                    lora_scale_1 = gr.Slider(label="LoRA 1 Scale", minimum=0, maximum=3, step=0.01, value=1.15)
+                    lora_scale_1 = gr.Slider(label="LoRA 1 Scale", minimum=0, maximum=3, step=0.05, value=1.15)
             with gr.Row():
                 remove_button_1 = gr.Button("Remove", size="sm")
         with gr.Column(scale=8):
@@ -510,7 +454,7 @@ with gr.Blocks(css=css, delete_cache=(60, 60)) as app:
                 with gr.Column(scale=3, min_width=100):
                     selected_info_2 = gr.Markdown("Select a LoRA 2")
                 with gr.Column(scale=5, min_width=50):
-                    lora_scale_2 = gr.Slider(label="LoRA 2 Scale", minimum=0, maximum=3, step=0.01, value=1.15)
+                    lora_scale_2 = gr.Slider(label="LoRA 2 Scale", minimum=0, maximum=3, step=0.05, value=1.15)
             with gr.Row():
                 remove_button_2 = gr.Button("Remove", size="sm")
     with gr.Row():
@@ -539,21 +483,16 @@ with gr.Blocks(css=css, delete_cache=(60, 60)) as app:
     with gr.Row():
         with gr.Accordion("Advanced Settings", open=False):
             with gr.Row():
-                input_image = gr.Image(label="Input image", type="filepath", show_share_button=False)
-                image_strength = gr.Slider(label="Denoise Strength", info="Lower means more image influence", minimum=0.1, maximum=1.0, step=0.01, value=0.75)
-            with gr.Column():
-                with gr.Row():
-                    cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.5)
-                    steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=28)
-
-                with gr.Row():
-                    width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=1024)
-                    height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=1024)
-
-                with gr.Row():
-                    randomize_seed = gr.Checkbox(True, label="Randomize seed")
-                    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
-                
+                cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=7.5)
+                steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=28)
+
+            with gr.Row():
+                width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=768)
+                height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=1024)
+
+            with gr.Row():
+                randomize_seed = gr.Checkbox(True, label="Randomize seed")
+                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
 
     gallery.select(
         update_selection,
@@ -588,7 +527,7 @@ with gr.Blocks(css=css, delete_cache=(60, 60)) as app:
     gr.on(
         triggers=[generate_button.click, prompt.submit],
         fn=run_lora,
-        inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, randomize_seed, seed, width, height, loras_state],
+        inputs=[prompt, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, randomize_seed, seed, width, height, loras_state],
         outputs=[result, seed, progress_bar]
     ).then(
         fn=lambda x, history: update_history(x, history),