Update app.py

app.py

@@ -1,3 +1,5 @@
+
+
 import os
 import uuid
 import gradio as gr
@@ -36,7 +38,125 @@ english_labels = {
     "Seed": "Seed"
 }
 
-# [Rest of the imports and pipeline setup remains the same...]
+# load pipelines
+base_model = "black-forest-labs/FLUX.1-schnell"
+
+taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.bfloat16).to("cuda")
+pipe = FluxPipeline.from_pretrained(base_model,
+                                    vae=taef1,
+                                    torch_dtype=torch.bfloat16)
+
+pipe.transformer.to(memory_format=torch.channels_last)
+clip_slider = CLIPSliderFlux(pipe, device=torch.device("cuda"))
+
+MAX_SEED = 2**32-1
+
+def save_images_with_unique_filenames(image_list, save_directory):
+    if not os.path.exists(save_directory):
+        os.makedirs(save_directory)
+
+    paths = []
+    for image in image_list:
+        unique_filename = f"{uuid.uuid4()}.png"
+        file_path = os.path.join(save_directory, unique_filename)
+        
+        image.save(file_path)
+        paths.append(file_path)
+    
+    return paths
+
+def convert_to_centered_scale(num):
+    if num % 2 == 0:  # even
+        start = -(num // 2 - 1)
+        end = num // 2
+    else:  # odd
+        start = -(num // 2)
+        end = num // 2 
+    return tuple(range(start, end + 1))
+
+def translate_if_korean(text):
+    if any('\u3131' <= char <= '\u3163' or '\uac00' <= char <= '\ud7a3' for char in text):
+        return translator(text)[0]['translation_text']
+    return text
+
+@spaces.GPU(duration=85)
+def generate(prompt,
+             concept_1,
+             concept_2,
+             scale,
+             randomize_seed=True,
+             seed=42,
+             recalc_directions=True,
+             iterations=200, 
+             steps=3, 
+             interm_steps=33, 
+             guidance_scale=3.5,
+             x_concept_1="", x_concept_2="", 
+             avg_diff_x=None, 
+             total_images=[],
+             progress=gr.Progress()
+    ):
+    # Translate prompt and concepts if Korean
+    prompt = translate_if_korean(prompt)
+    concept_1 = translate_if_korean(concept_1)
+    concept_2 = translate_if_korean(concept_2)
+    
+    print(f"Prompt: {prompt}, ← {concept_2}, {concept_1} ➡️ . scale {scale}, interm steps {interm_steps}")
+    slider_x = [concept_2, concept_1]
+    # check if avg diff for directions need to be re-calculated
+    if randomize_seed:
+            seed = random.randint(0, MAX_SEED)
+        
+    if not sorted(slider_x) == sorted([x_concept_1, x_concept_2]) or recalc_directions:
+        progress(0, desc="Calculating directions...")
+        avg_diff = clip_slider.find_latent_direction(slider_x[0], slider_x[1], num_iterations=iterations)
+        x_concept_1, x_concept_2 = slider_x[0], slider_x[1]
+
+    images = []
+    high_scale = scale
+    low_scale = -1 * scale
+    for i in progress.tqdm(range(interm_steps), desc="Generating images"):
+        cur_scale = low_scale + (high_scale - low_scale) * i / (interm_steps - 1)
+        image = clip_slider.generate(prompt, 
+                                     width=768,
+                                     height=768,
+                                     guidance_scale=guidance_scale, 
+                                     scale=cur_scale,  seed=seed, num_inference_steps=steps, avg_diff=avg_diff) 
+        images.append(image)
+    canvas = Image.new('RGB', (256*interm_steps, 256))
+    for i, im in enumerate(images):
+        canvas.paste(im.resize((256,256)), (256 * i, 0))
+
+    comma_concepts_x = f"{slider_x[1]}, {slider_x[0]}"
+
+    scale_total = convert_to_centered_scale(interm_steps)
+    scale_min = scale_total[0]
+    scale_max = scale_total[-1]
+    scale_middle = scale_total.index(0)
+    post_generation_slider_update = gr.update(label=comma_concepts_x, value=0, minimum=scale_min, maximum=scale_max, interactive=True)
+    avg_diff_x = avg_diff.cpu()
+
+    video_path = f"{uuid.uuid4()}.mp4"
+    print(video_path)
+    return x_concept_1,x_concept_2, avg_diff_x, export_to_video(images, video_path, fps=5), canvas, images, images[scale_middle], post_generation_slider_update, seed
+
+def update_pre_generated_images(slider_value, total_images):
+    number_images = len(total_images)
+    if(number_images > 0):
+        scale_tuple = convert_to_centered_scale(number_images)
+        return total_images[scale_tuple.index(slider_value)][0]
+    else:
+        return None
+    
+def reset_recalc_directions():
+    return True
+
+# Updated examples with English text
+examples = [
+    ["flower in mountain", "spring", "winter", 1.5],
+    ["man", "baby", "elderly", 2.5],
+    ["a tomato", "super fresh", "rotten", 2.5]
+]
 
 css = """
 footer {
@@ -127,13 +247,6 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
                                interactive=True, 
                                randomize=True)
 
-    # Updated examples with English text
-    examples = [
-        ["flower in mountain", "spring", "winter", 1.5],
-        ["man", "baby", "elderly", 2.5],
-        ["a tomato", "super fresh", "rotten", 2.5]
-    ]
-
     examples_gradio = gr.Examples(
         examples=examples,
         inputs=[prompt, concept_1, concept_2, x],
@@ -143,7 +256,30 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
         cache_examples="lazy"
     )
 
-    # [Rest of the event handlers remain the same...]
-
+    submit.click(
+        fn=generate,
+        inputs=[prompt, concept_1, concept_2, x, randomize_seed, seed, recalc_directions, 
+                iterations, steps, interm_steps, guidance_scale, x_concept_1, x_concept_2, 
+                avg_diff_x, total_images],
+        outputs=[x_concept_1, x_concept_2, avg_diff_x, output_image, image_seq, total_images, 
+                post_generation_image, post_generation_slider, seed]
+    )
+    iterations.change(
+        fn=reset_recalc_directions,
+        outputs=[recalc_directions]
+    )
+    seed.change(
+        fn=reset_recalc_directions,
+        outputs=[recalc_directions]
+    )
+    post_generation_slider.change(
+        fn=update_pre_generated_images,
+        inputs=[post_generation_slider, total_images],
+        outputs=[post_generation_image],
+        queue=False,
+        show_progress="hidden",
+        concurrency_limit=None
+    )
+        
 if __name__ == "__main__":
     demo.launch()