chore: Refactor image enhancement functionality and update requirements for diffusers

app.py

@@ -1,85 +1,36 @@
+import torch
+from diffusers import FluxPriorReduxPipeline, FluxPipeline
+from diffusers.utils import load_image
 import gradio as gr
 import spaces
-import torch
-from gradio_client import Client, handle_file
-from colorama import Fore, Style
-from diffusers import AutoPipelineForImage2Image
-from PIL import Image
-from live_preview_helpers import flux_pipe_call_that_returns_an_iterable_of_images
-
-joy_client = Client("fancyfeast/joy-caption-alpha-two")
-qwen_client = Client("Qwen/Qwen2.5-72B-Instruct")
-
-pipeline = AutoPipelineForImage2Image.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
-pipeline.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipeline)
-
-lora_ids = {
-    # "Realism": "XLabs-AI/flux-RealismLora",
-    "Cartoonism": "aleksa-codes/flux-ghibsky-illustration",
-}
-
-def load_lora(lora_name):
-    print(f"Loading LoRA model: {lora_name}")
-    global pipeline
-    pipeline.unload_lora_weights()
-    pipeline.load_lora_weights(lora_ids[lora_name])
-    pipeline.enable_model_cpu_offload()
-    print(f"{Fore.GREEN}LoRA model loaded{Style.RESET_ALL}")
-
-def describe_image(image_path):
-    print(f"Describing image: {image_path}")
-    image_description = joy_client.predict(
-            input_image=handle_file(image_path),
-            caption_type="Descriptive",
-            caption_length="long",
-            extra_options=[],
-            name_input="",
-            custom_prompt="",
-            api_name="/stream_chat"
-    )[-1]
-    print(f"{Fore.GREEN}{image_description}{Style.RESET_ALL}")
-    return image_description
-
-def refine_prompt(image_description):
-    print(f"Improving prompt: {image_description}")
-    qwen_prompt = f"""This is the description of the image: {image_description}
-
-And those some good AI Art Prompts:
-- a cat on a windowsill gazing out at a starry night sky and distant city lights
-- a fisherman casting a line into a peaceful village lake surrounded by quaint cottages
-- cozy mountain cabin covered in snow, with smoke curling from the chimney and a warm, inviting light spilling through the windows
-- Mykonos
-- an orange Lamborghini driving down a hill road at night with a beautiful ocean view in the background, side view, no text
-- a small Yorkie on a windowsill during a snowy winter night, with a warm, cozy glow from inside and soft snowflakes drifting outside
-- serene Japanese garden with a koi pond and a traditional tea house, nestled under a canopy of cherry blossoms in full bloom
-- the most beautiful place in the universe
-
-Based on what i gave you, Write a great short AI Art Prompt for me that is based on the image description above (Don't write anything else, just the prompt)
-"""
-    refined_prompt = qwen_client.predict(
-            query=qwen_prompt,
-            history=[],
-            system="You are Qwen, created by Alibaba Cloud. You are a helpful assistant.",
-            api_name="/model_chat"
-    )[1][0][-1]
-    print(f"{Fore.GREEN}{refined_prompt}{Style.RESET_ALL}")
-    return refined_prompt
-
-@spaces.GPU(duration=75)
-def img2img_infer(image_path, image_description):
-    pil_image = Image.open(image_path)
-    width, height = pil_image.size
-    for enhanced_image in pipeline.flux_pipe_call_that_returns_an_iterable_of_images(
-        prompt=f'GHIBSKY style, {image_description}',
-        image=pil_image,
-        # guidance_scale=3.5,
-        # num_inference_steps=28,
-        # width=1024,
-        # height=1024,
-        # generator=torch.Generator("cpu").manual_seed(0),
-        output_type="pil",
-    ):
-        yield enhanced_image.resize((width, height))
+# from live_preview_helpers import flux_pipe_call_that_returns_an_iterable_of_images
+
+pipe_prior_redux = FluxPriorReduxPipeline.from_pretrained("black-forest-labs/FLUX.1-Redux-dev", revision="refs/pr/8", torch_dtype=torch.bfloat16).to("cuda")
+pipe = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-dev" , 
+    text_encoder=None,
+    text_encoder_2=None,
+    torch_dtype=torch.bfloat16
+)
+# pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
+# pipe.enable_sequential_cpu_offload()
+
+@spaces.GPU(duration=120)
+def enhance_image(image_path, keep_aspect_ratio=False):
+    print(image_path)
+    image = load_image(image_path)
+    print(image.size)
+    width, height = image.size if keep_aspect_ratio else (None, None)
+    pipe_prior_output = pipe_prior_redux(image)
+    images = pipe(
+        height=height,
+        width=width,
+        guidance_scale=2.5,
+        num_inference_steps=50,
+        generator=torch.Generator("cpu").manual_seed(0),
+        **pipe_prior_output,
+    ).images
+    return images[0]
 
 
 with gr.Blocks(title="Magnific") as demo:
@@ -89,17 +40,10 @@ with gr.Blocks(title="Magnific") as demo:
     with gr.Row():
         with gr.Column():
             image_path = gr.Image(label="Image", type="filepath")
-            lora_dropdown = gr.Dropdown(label="LoRA Model", choices=list(lora_ids.keys()), value=None)
-            describe_btn = gr.Button(value="Describe Image", variant="primary")
-            with gr.Row(equal_height=True):
-                image_description = gr.Textbox(label="Image Description", scale=4)
-                refine_prompt_btn = gr.Button(value="Refine", variant="primary", scale=1)
+            keep_aspect_ratio = gr.Checkbox(label="Keep Aspect Ratio", value=False)
             submit_btn = gr.Button(value="Submit", variant="primary")
         enhanced_image = gr.Image(label="Enhanced Image", type="pil")
 
-    lora_dropdown.change(load_lora, inputs=lora_dropdown)
-    refine_prompt_btn.click(refine_prompt, inputs=image_description, outputs=image_description)
-    describe_btn.click(describe_image, inputs=image_path, outputs=image_description)
-    submit_btn.click(img2img_infer, inputs=[image_path, image_description], outputs=enhanced_image)
+    submit_btn.click(enhance_image, inputs=[image_path, keep_aspect_ratio], outputs=enhanced_image)
 
 demo.queue().launch(share=False)

requirements.txt

@@ -1,6 +1,6 @@
 colorama
 gradio-client
-diffusers
+git+https://github.com/huggingface/diffusers.git
 torch
 transformers
 spaces