Add i2i

app.py

@@ -5,9 +5,9 @@ import logging
 import torch
 from PIL import Image
 import spaces
-from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL
+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
 import copy
 import random
@@ -25,6 +25,15 @@ base_model = "black-forest-labs/FLUX.1-dev"
 taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
 good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype).to(device)
 pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1).to(device)
+pipe_i2i = AutoPipelineForImage2Image.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
+                                                     )
 
 MAX_SEED = 2**32-1
 
@@ -88,7 +97,26 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scal
         ):
             yield img
 
-def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
+@spaces.GPU(duration=70)
+def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps, cfg_scale, width, height, lora_scale, seed):
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+    pipe_i2i.to("cuda")
+    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 
+    
+def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
     if selected_index is None:
         raise gr.Error("You must select a LoRA before proceeding.")
     selected_lora = loras[selected_index]
@@ -107,32 +135,44 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, wid
 
     with calculateDuration("Unloading LoRA"):
         pipe.unload_lora_weights()
+        pipe_i2i.unload_lora_weights()
         
     # Load LoRA weights
     with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"):
-        if "weights" in selected_lora:
-            pipe.load_lora_weights(lora_path, weight_name=selected_lora["weights"])
+        if(image_input is not None):
+            if "weights" in selected_lora:
+                pipe_i2i.load_lora_weights(lora_path, weight_name=selected_lora["weights"])
+            else:
+                pipe_i2i.load_lora_weights(lora_path)
         else:
-            pipe.load_lora_weights(lora_path)
-
+            if "weights" in selected_lora:
+                pipe.load_lora_weights(lora_path, weight_name=selected_lora["weights"])
+            else:
+                pipe.load_lora_weights(lora_path)
+                
     # Set random seed for reproducibility
     with calculateDuration("Randomizing seed"):
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
-
-    image_generator = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, progress)
+            
+    if(image_input is not None):
+        
+        final_image = generate_image_to_image(prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, lora_scale, seed)
+        yield final_image, seed, gr.update(visible=False)
+    else:
+        image_generator = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, 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)
+        # 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(value=progress_bar, visible=False)
-
 def get_huggingface_safetensors(link):
   split_link = link.split("/")
   if(len(split_link) == 2):
@@ -257,6 +297,9 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css, delete_cache=(60, 3600)) as app:
 
     with gr.Row():
         with gr.Accordion("Advanced Settings", open=False):
+            with gr.Row():
+                input_image = gr.Image(label="Input image", type="filepath")
+                image_strength = gr.Slider(label="Image Strength", 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)
@@ -288,7 +331,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css, delete_cache=(60, 3600)) as app:
     gr.on(
         triggers=[generate_button.click, prompt.submit],
         fn=run_lora,
-        inputs=[prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale],
+        inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale],
         outputs=[result, seed, progress_bar]
     )