Update app.py

app.py

@@ -6,8 +6,10 @@ import spaces
 import torch
 import json
 import logging
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL, AutoPipelineForImage2Image
 from huggingface_hub import login
+from diffusers.utils import load_image
+
 import time
 from datetime import datetime
 from io import BytesIO
@@ -23,7 +25,6 @@ import json
 HF_TOKEN = os.environ.get("HF_TOKEN")
 login(token=HF_TOKEN)
 import diffusers
-print(diffusers.__version__)
 
 # init
 dtype = torch.float16  # use float16 for fast generate
@@ -31,8 +32,23 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 base_model = "black-forest-labs/FLUX.1-dev"
 
 # load pipe
+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).to(device)
 
+# img2img model
+img2img = 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
 
 class calculateDuration:
@@ -56,8 +72,7 @@ class calculateDuration:
             print(f"Elapsed time: {self.elapsed_time:.6f} seconds")
         
 @spaces.GPU(duration=120)
-@torch.inference_mode()
-def generate_image(prompt, adapter_names,  steps, seed, cfg_scale, width, height, progress):
+def generate_image(orginal_image, prompt, adapter_names,  steps, seed, image_strength, cfg_scale, width, height, progress):
 
     
     gr.Info("Start to generate images ...")
@@ -67,15 +82,28 @@ def generate_image(prompt, adapter_names,  steps, seed, cfg_scale, width, height
     
     with calculateDuration("Generating image"):
         # Generate image
-        generated_image = pipe(
-            prompt=prompt,
-            num_inference_steps=steps,
-            guidance_scale=cfg_scale,
-            width=width,
-            height=height,
-            max_sequence_length=512,
-            generator=generator,
-        ).images[0]
+        if orginal_image:
+            generated_image = img2img(
+                prompt=prompt,
+                image=orginal_image,
+                strength=image_strength,
+                num_inference_steps=steps,
+                guidance_scale=cfg_scale,
+                width=width,
+                height=height,
+                generator=generator,
+                joint_attention_kwargs={"scale": lora_scale}
+            ).images[0]
+        else:
+            generated_image = pipe(
+                prompt=prompt,
+                num_inference_steps=steps,
+                guidance_scale=cfg_scale,
+                width=width,
+                height=height,
+                max_sequence_length=512,
+                generator=generator,
+            ).images[0]
     
     progress(99, "Generate image success!")
     return generated_image
@@ -119,10 +147,15 @@ def upload_image_to_r2(image, account_id, access_key, secret_key, bucket_name):
         print("upload thumbnail finish", thumbnail_file)
     return image_file
 
-def run_lora(prompt, lora_strings_json,  cfg_scale, steps, randomize_seed, seed, width, height, upload_to_r2, account_id, access_key, secret_key, bucket, progress=gr.Progress(track_tqdm=True)):
+def run_lora(prompt, image_url, lora_strings_json, image_strength,  cfg_scale, steps, randomize_seed, seed, width, height, upload_to_r2, account_id, access_key, secret_key, bucket, progress=gr.Progress(track_tqdm=True)):
     print("run_lora", prompt, lora_strings_json, cfg_scale, steps, width, height)
     gr.Info("Starting process")
-    
+
+    img2img_model = False
+    orginal_image = None
+    if image_url:
+        orginal_image = load_image(image_url)
+        img2img_model = True
     # Set random seed for reproducibility
     if randomize_seed:
         with calculateDuration("Set random seed"):
@@ -152,7 +185,10 @@ def run_lora(prompt, lora_strings_json,  cfg_scale, steps, randomize_seed, seed,
                         retry_count = 3
                         for attempt in range(retry_count):
                             try:
-                                pipe.load_lora_weights(lora_repo, weight_name=weights, adapter_name=adapter_name)
+                                if img2img_model:
+                                    img2img.load_lora_weights(lora_repo, weight_name=weights, adapter_name=adapter_name)
+                                else:
+                                    pipe.load_lora_weights(lora_repo, weight_name=weights, adapter_name=adapter_name)
                                 adapter_names.append(adapter_name)
                                 adapter_weights.append(adapter_weight)
                                 break  # Load successful, exit retry loop
@@ -165,14 +201,17 @@ def run_lora(prompt, lora_strings_json,  cfg_scale, steps, randomize_seed, seed,
                         
                 # set lora weights
                 if len(adapter_names) > 0:
-                    pipe.set_adapters(adapter_names, adapter_weights=adapter_weights)
+                    if img2img_model:
+                        img2img.set_adapters(adapter_names, adapter_weights=adapter_weights)
+                    else:
+                        pipe.set_adapters(adapter_names, adapter_weights=adapter_weights)
 
     
     # Generate image
     error_message = ""
     try:
         print("Start applying for zeroGPU resources")
-        final_image = generate_image(prompt, adapter_names, steps, seed, cfg_scale, width, height, progress)
+        final_image = generate_image(orginal_image, prompt, adapter_names, steps, seed, image_strength, cfg_scale, width, height, progress)
     except Exception as e:
         error_message =  str(e)
         gr.Error(error_message)
@@ -210,7 +249,7 @@ with gr.Blocks(css=css) as demo:
 
             prompt = gr.Text(label="Prompt", placeholder="Enter prompt", lines=10)
             lora_strings_json = gr.Text(label="LoRA Configs (JSON List String)", placeholder='[{"repo": "lora_repo1", "weights": "weights1", "adapter_name": "adapter_name1", "adapter_weight": 1}, {"repo": "lora_repo2", "weights": "weights2", "adapter_name": "adapter_name2", "adapter_weight": 1}]', lines=5)
-            
+            image_url = gr.Text(label="Image url", placeholder="Enter image url to enable image to image model", lines=1)
             run_button = gr.Button("Run", scale=0)
 
             with gr.Accordion("Advanced Settings", open=False):
@@ -224,6 +263,7 @@ with gr.Blocks(css=css) as demo:
                     height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=1024)
 
                 with gr.Row():
+                    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)
                     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) 
 
@@ -244,7 +284,9 @@ with gr.Blocks(css=css) as demo:
     )
     inputs = [
         prompt,
+        image_url,
         lora_strings_json,
+        image_strength,
         cfg_scale,
         steps,
         randomize_seed,