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]
  )