initial commit

app.py

@@ -0,0 +1,147 @@
+import torch
+import numpy as np
+import math
+import spaces
+from diffusers import StableDiffusionXLPipeline
+from transformers import AutoFeatureExtractor
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
+from ip_adapter.ip_adapter import EasyRef
+from huggingface_hub import hf_hub_download
+import gradio as gr
+import cv2
+import pillow_avif
+
+def adaptive_resize(w, h, size=1024):
+    times = math.sqrt(h * w / (size**2))
+    if w==h:
+        w, h = size, size
+    elif times > 1.1:
+        w, h = math.ceil(w / times), math.ceil(h / times)
+    elif times < 0.8:
+        w, h = math.ceil(w / times), math.ceil(h / times)
+    new_w, new_h = 64 * (math.ceil(w / 64)), 64 * (math.ceil(h / 64))
+    return new_w, new_h
+
+def res2string(w, h):
+    return str(w)+"x"+str(h)
+
+base_model_path = "stabilityai/stable-diffusion-xl-base-1.0"
+multimodal_llm_path = "Qwen/Qwen2-VL-2B-Instruct"
+ip_ckpt = hf_hub_download(repo_id="zongzhuofan/EasyRef", filename="pytorch_model.bin", repo_type="model")
+
+safety_model_id = "CompVis/stable-diffusion-safety-checker"
+safety_feature_extractor = AutoFeatureExtractor.from_pretrained(safety_model_id)
+safety_checker = StableDiffusionSafetyChecker.from_pretrained(safety_model_id)
+
+device = "cuda"
+
+pipe = StableDiffusionXLPipeline.from_pretrained(
+    base_model_path,
+    torch_dtype=torch.float16,
+    feature_extractor=safety_feature_extractor,
+    safety_checker=safety_checker,
+    add_watermarker=False,
+).to(device)
+
+easyref = EasyRef(pipe, multimodal_llm_path, ip_ckpt, device, num_tokens=64, use_lora=True, lora_rank=128)
+
+cv2.setNumThreads(1)
+
+@spaces.GPU(enable_queue=True)
+def generate_image(images, prompt, negative_prompt, height, width, scale, num_inference_steps, seed, progress=gr.Progress(track_tqdm=True)):
+    print("Generating")
+    new_width, new_height = adaptive_resize(width, height)
+    print("The output resolution is adaptively resized from {ori} to {new}".format(ori=res2string(width, height), new=res2string(new_width, new_height)))
+    template = "Visualize a scene that closely resembles the provided images, capturing the essence and details described in this prompt:\n"
+    system_prompt = [template + prompt, template]
+    image = easyref.generate(
+        pil_image=images,
+        system_prompt=system_prompt, 
+        prompt=prompt, 
+        negative_prompt=negative_prompt, 
+        num_samples=1,
+        num_inference_steps=num_inference_steps,
+        seed=seed,
+        height=new_height,
+        width=new_width)    
+    print(image)
+    return image
+
+# def change_style(style):
+#     if style == "Photorealistic":
+#         return(gr.update(value=True), gr.update(value=1.3), gr.update(value=1.0))
+#     else:
+#         return(gr.update(value=True), gr.update(value=0.1), gr.update(value=0.8))
+
+def swap_to_gallery(images):
+    return gr.update(value=images, visible=True), gr.update(visible=True), gr.update(visible=False)
+
+def remove_back_to_files():
+    return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
+
+MAX_SEED = np.iinfo(np.int32).max
+css = '''
+h1{margin-bottom: 0 !important}
+'''
+with gr.Blocks(css=css) as demo:
+    gr.Markdown("# EasyRef demo")
+    gr.Markdown("Demo for the [zongzhuofan/EasyRef model](https://huggingface.co/zongzhuofan/EasyRef)")
+    with gr.Row():
+        with gr.Column():
+            files = gr.Files(
+                        label="Upload several reference images",
+                        file_types=["image"]
+                    )
+            uploaded_files = gr.Gallery(label="Your images", visible=False, columns=5, rows=1, height=125)
+            with gr.Column(visible=False) as clear_button:
+                remove_and_reupload = gr.ClearButton(value="Remove and upload new ones", components=files, size="sm")
+            prompt = gr.Textbox(label="Prompt",
+                       placeholder="An oil painting of a [man/woman/person]...")
+            negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="A collage of images, naked, monochrome, lowres, bad anatomy, worst quality, low quality")
+            # style = gr.Radio(label="Generation type", info="For stylized try prompts like 'a watercolor painting of a woman'", choices=["Photorealistic", "Stylized"], value="Photorealistic")
+            submit = gr.Button("Submit")
+            with gr.Accordion(open=False, label="Advanced Options"):
+                height = gr.Slider(label="Height", info="This will be adaptively resized", value=1024, step=64, minimum=896, maximum=1280)
+                width = gr.Slider(label="Width", info="This will be adaptively resized", value=1024, step=64, minimum=896, maximum=1280)
+                scale = gr.Slider(label="Scale", info="Scale for reference images", value=1.0, step=0.1, minimum=0.5, maximum=1.5)
+                num_inference_steps = gr.Slider(label="Number of inference steps", value=30, step=1, minimum=1, maximum=60)
+                seed = gr.Slider(label="Seed", value=24, step=1, minimum=0, maximum=MAX_SEED)
+        with gr.Column():
+            gallery = gr.Gallery(label="Generated Images")
+        # style.change(fn=change_style,
+        #             inputs=style,
+        #             outputs=[preserve, face_strength, likeness_strength])
+        files.upload(fn=swap_to_gallery, inputs=files, outputs=[uploaded_files, clear_button, files])
+        remove_and_reupload.click(fn=remove_back_to_files, outputs=[uploaded_files, clear_button, files])
+        submit.click(fn=generate_image,
+                    inputs=[files,prompt,negative_prompt,height, width, scale, num_inference_steps, seed],
+                    outputs=gallery)
+        examples = gr.Examples(
+            examples=[
+                [
+                    [
+                        "assets/aragaki_identity/1.jpg",
+                        "assets/aragaki_identity/2.webp",
+                        "assets/aragaki_identity/3.webp",
+                        "assets/aragaki_identity/4.jpeg",
+                        "assets/aragaki_identity/5.webp",
+                    ],
+                    "An oil painting of a smiling woman.",
+                ],
+                [
+                    [
+                        "assets/blindbox_style/1.jpg",
+                        "assets/blindbox_style/2.jpg",
+                        "assets/blindbox_style/3.jpg",
+                        "assets/blindbox_style/4.jpg",
+                        "assets/blindbox_style/5.jpg",
+                    ],
+                    "Donald Trump"
+                ],
+            ],
+            inputs=[files, prompt],
+        )
+            
+    gr.Markdown("We release our checkpoints for research purposes only. Users are granted the freedom to create images using this tool, but they are expected to comply with local laws and utilize it in a responsible manner. The developers do not assume any responsibility for potential misuse by users.")
+    
+demo.launch()

requirements.txt

@@ -0,0 +1,14 @@
+diffusers==0.22.1
+torch==2.3.1
+torchvision==0.18.1
+transformers==4.45.2
+qwen_vl_utils==0.0.8
+accelerate
+safetensors
+einops
+omegaconf
+peft
+pillow-avif-plugin
+huggingface-hub==0.23.3
+opencv-python
+git+https://github.com/TempleX98/EasyRef.git