Update app.py

app.py

@@ -13,7 +13,9 @@ import gradio as gr
 import spaces
 from diffusers import (
     DiffusionPipeline,
-    FlowMatchEulerDiscreteScheduler)
+    FlowMatchEulerDiscreteScheduler,
+    AutoencoderKL,
+    AutoPipelineForImage2Image)
 from huggingface_hub import (
     hf_hub_download,
     HfFileSystem,
@@ -142,15 +144,30 @@ pipe = DiffusionPipeline.from_pretrained(
     base_model, scheduler=scheduler, torch_dtype=dtype
 ).to(device)
 
+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.vae = taef1
+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,
+    scheduler=scheduler,
+    torch_dtype=dtype
+).to(device)
+
 # Lightning LoRA info (no global state)
 LIGHTNING_LORA_REPO = "lightx2v/Qwen-Image-Lightning"
 LIGHTNING_LORA_WEIGHT = "Qwen-Image-Lightning-8steps-V1.0.safetensors"
 
-MAX_SEED = np.iinfo(np.int32).max
+MAX_SEED = 2**32 - 1
 
-class Timer:
-    def __init__(self, task_name=""):
-        self.task_name = task_name
+class calculateDuration:
+    def __init__(self, activity_name=""):
+        self.activity_name = activity_name
 
     def __enter__(self):
         self.start_time = time.time()
@@ -159,8 +176,8 @@ class Timer:
     def __exit__(self, exc_type, exc_value, traceback):
         self.end_time = time.time()
         self.elapsed_time = self.end_time - self.start_time
-        if self.task_name:
-            print(f"Elapsed time for {self.task_name}: {self.elapsed_time:.6f} seconds")
+        if self.activity_name:
+            print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds")
         else:
             print(f"Elapsed time: {self.elapsed_time:.6f} seconds")
 
@@ -213,26 +230,85 @@ def adjust_generation_mode(speed_mode):
         return gr.update(value="Base mode selected - 48 steps for best quality"), 48, 4.0
 
 @spaces.GPU(duration=100)
-def create_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, negative_prompt=""):
+def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, negative_prompt=""):
+    generator = torch.Generator(device="cuda").manual_seed(seed)
     pipe.to("cuda")
+
+    batch_size = 1
+    prompt = prompt_mash
+    do_classifier_free_guidance = cfg_scale > 1.0
+    prompt_embeds, pooled_prompt_embeds = pipe.encode_prompt(
+        prompt,
+        num_images_per_prompt=1,
+        do_classifier_free_guidance=do_classifier_free_guidance,
+        prompt_2=None,
+        max_sequence_length=256,
+    )
+    height, width = height - height % 16, width - width % 16
+    latents = pipe.prepare_latents(
+        batch_size,
+        pipe.transformer.config.in_channels,
+        height,
+        width,
+        dtype,
+        device,
+        generator,
+        latents=None,
+    )
+    pipe.scheduler.set_timesteps(steps)
+    timesteps = pipe.scheduler.timesteps
+    joint_attention_kwargs = {"scale": lora_scale}
+    for i in range(steps):
+        t = pipe.scheduler.sigmas[i]
+        latent_model_input = latents
+        with torch.no_grad():
+            noise_pred = pipe.transformer(
+                hidden_states=latent_model_input,
+                timestep=t,
+                guidance=cfg_scale,
+                pooled_projections=pooled_prompt_embeds,
+                encoder_hidden_states=prompt_embeds,
+                joint_attention_kwargs=joint_attention_kwargs,
+                return_dict=False,
+            )[0]
+        latents = pipe.scheduler.step(
+            model_output=noise_pred,
+            timestep=t,
+            sample=latent_model_input,
+            return_dict=False,
+        )[0]
+        # preview
+        with torch.no_grad():
+            decoded = pipe.vae.decode(latents / pipe.vae.config.scaling_factor, return_dict=False)[0]
+        image = pipe.image_processor.pt_to_pil(decoded)[0]
+        yield image
+    # final
+    with torch.no_grad():
+        decoded = good_vae.decode(latents / good_vae.config.scaling_factor, return_dict=False)[0]
+    image = pipe.image_processor.pt_to_pil(decoded)[0]
+    yield image
+
+@spaces.GPU(duration=100)
+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)
-    
-    with Timer("Generating image"):
-        # Generate image
-        image = pipe(
-            prompt=prompt_mash,
-            negative_prompt=negative_prompt,
-            num_inference_steps=steps,
-            true_cfg_scale=cfg_scale,  # Use true_cfg_scale for Qwen-Image
-            width=width,
-            height=height,
-            generator=generator,
-        ).images[0]
-        
-    return image
+    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
 
 @spaces.GPU(duration=100)
-def process_adapter_generation(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, aspect_ratio, lora_scale, speed_mode, progress=gr.Progress(track_tqdm=True)):
+def process_adapter_generation(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, aspect_ratio, lora_scale, speed_mode, image_input, image_strength, negative_prompt="", progress=gr.Progress(track_tqdm=True)):
     if selected_index is None:
         raise gr.Error("You must select a LoRA before proceeding.")
     
@@ -253,14 +329,16 @@ def process_adapter_generation(prompt, cfg_scale, steps, selected_index, randomi
         prompt_mash = prompt
 
     # Always unload any existing LoRAs first to avoid conflicts
-    with Timer("Unloading existing LoRAs"):
+    with calculateDuration("Unloading existing LoRAs"):
         pipe.unload_lora_weights()
+        pipe_i2i.unload_lora_weights()
+
+    pipe_to_use = pipe_i2i if image_input is not None else pipe
 
-    # Load LoRAs based on speed mode
     if speed_mode == "Fast (8 steps)":
-        with Timer("Loading Lightning LoRA and style LoRA"):
+        with calculateDuration("Loading Lightning LoRA and style LoRA"):
             # Load Lightning LoRA first
-            pipe.load_lora_weights(
+            pipe_to_use.load_lora_weights(
                 LIGHTNING_LORA_REPO, 
                 weight_name=LIGHTNING_LORA_WEIGHT,
                 adapter_name="lightning"
@@ -268,7 +346,7 @@ def process_adapter_generation(prompt, cfg_scale, steps, selected_index, randomi
             
             # Load the selected style LoRA
             weight_name = selected_lora.get("weights", None)
-            pipe.load_lora_weights(
+            pipe_to_use.load_lora_weights(
                 lora_path, 
                 weight_name=weight_name, 
                 low_cpu_mem_usage=True,
@@ -276,29 +354,36 @@ def process_adapter_generation(prompt, cfg_scale, steps, selected_index, randomi
             )
             
             # Set both adapters active with their weights
-            pipe.set_adapters(["lightning", "style"], adapter_weights=[1.0, lora_scale])
+            pipe_to_use.set_adapters(["lightning", "style"], adapter_weights=[1.0, lora_scale])
     else:
         # Quality mode - only load the style LoRA
-        with Timer(f"Loading LoRA weights for {selected_lora['title']}"):
+        with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"):
             weight_name = selected_lora.get("weights", None)
-            pipe.load_lora_weights(
+            pipe_to_use.load_lora_weights(
                 lora_path, 
                 weight_name=weight_name, 
                 low_cpu_mem_usage=True
             )
                 
     # Set random seed for reproducibility
-    with Timer("Randomizing seed"):
+    with calculateDuration("Randomizing seed"):
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
     
     # Get image dimensions from aspect ratio
     width, height = compute_image_dimensions(aspect_ratio)
     
-    # Generate the image
-    final_image = create_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale)
+    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, negative_prompt)
     
-    return final_image, seed
+        step_counter = 0
+        for image in image_generator:
+            step_counter += 1
+            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)
 
 def fetch_hf_adapter_files(link):
     split_link = link.split("/")
@@ -422,8 +507,6 @@ def incorporate_custom_adapter(custom_lora):
 def discard_custom_adapter():
     return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None, ""
 
-process_adapter_generation.zerogpu = True
-
 css = '''
 #gen_btn{height: 100%}
 #gen_column{align-self: stretch}
@@ -436,6 +519,10 @@ css = '''
 .card_internal img{margin-right: 1em}
 .styler{--form-gap-width: 0px !important}
 #speed_status{padding: .5em; border-radius: 5px; margin: 1em 0}
+#progress{height:30px}
+#progress .generating{display:none}
+.progress-container {width: 100%;height: 30px;background-color: #f0f0f0;border-radius: 15px;overflow: hidden;margin-bottom: 20px}
+.progress-bar {height: 100%;background-color: #4f46e5;width: calc(var(--current) / var(--total) * 100%);transition: width 0.5s ease-in-out}
 '''
 
 with gr.Blocks(theme="bethecloud/storj_theme", css=css, delete_cache=(120, 120)) as app:
@@ -467,6 +554,7 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css, delete_cache=(120, 120))
         
         with gr.Column():
             result = gr.Image(label="Generated Image")
+            progress_html = gr.HTML(visible=False, elem_id="progress")
 
             with gr.Row():
                 aspect_ratio = gr.Dropdown(
@@ -508,6 +596,10 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css, delete_cache=(120, 120))
                     randomize_seed = gr.Checkbox(True, label="Randomize seed")
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
                     lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=2, step=0.01, value=1.0)
+                
+                with gr.Row():
+                    image_input = gr.Image(label="Input Image for Image2Image", type="filepath")
+                    image_strength = gr.Slider(label="Image Strength", minimum=0, maximum=1, step=0.01, value=0.35)
 
     # Event handlers
     gallery.select(
@@ -536,8 +628,8 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css, delete_cache=(120, 120))
     gr.on(
         triggers=[generate_button.click, prompt.submit],
         fn=process_adapter_generation,
-        inputs=[prompt, cfg_scale, steps, selected_index, randomize_seed, seed, aspect_ratio, lora_scale, speed_mode],
-        outputs=[result, seed]
+        inputs=[prompt, cfg_scale, steps, selected_index, randomize_seed, seed, aspect_ratio, lora_scale, speed_mode, image_input, image_strength],
+        outputs=[result, seed, progress_html]
     )
 
 app.queue()