init

.gitignore

@@ -0,0 +1,5 @@
+__pycache__
+
+.env.*
+
+.DS_Store

.python-version

@@ -0,0 +1 @@
+3.11

adapter.py

@@ -0,0 +1,44 @@
+import os
+import yaml
+
+from huggingface_hub import hf_hub_download
+
+from src.models.sdxl.adapter.ip_adapter import (
+    SDXLModelWithIPAdapter,
+    SDXLModelWithIPAdapterConfig,
+)
+
+
+def _load_config(config_path: str, model_path: str, adapter_path: str):
+    with open(config_path, "r") as f:
+        config = yaml.safe_load(f)
+
+    config = SDXLModelWithIPAdapterConfig(**config)
+    config.checkpoint_path = model_path
+    config.adapter.checkpoint_weight = adapter_path
+
+    return config
+
+
+def load_ip_adapter_model(
+    model_path: str,
+    config_path: str,
+    adapter_path: str,
+):
+    config = _load_config(config_path, model_path, adapter_path)
+    model = SDXLModelWithIPAdapter.from_checkpoint(config)
+
+    return model
+
+
+def get_file_path(repo_id: str, filename: str, revision: str = "main") -> str:
+    if os.path.exists(filename):
+        return filename
+
+    file_path = hf_hub_download(
+        repo_id=repo_id,
+        filename=filename,
+        revision=revision,
+    )
+
+    return file_path

app.py

@@ -1,154 +1,219 @@
-import gradio as gr
-import numpy as np
+import spaces
+
+import os
 import random
+from PIL import Image
 
-# import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
 import torch
+import gradio as gr
+import dotenv
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
+from adapter import load_ip_adapter_model, get_file_path
+from example import EXAMPLES
 
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
+dotenv.load_dotenv(".env.local")
 
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
+ADAPTER_REPO_ID = os.environ.get("ADAPTER_REPO_ID")
+ADAPTER_MODEL_PATH = os.environ.get("ADAPTER_MODEL_PATH")
+ADAPTER_CONFIG_PATH = os.environ.get("ADAPTER_CONFIG_PATH")
 
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
+assert ADAPTER_REPO_ID is not None
+assert ADAPTER_MODEL_PATH is not None
+assert ADAPTER_CONFIG_PATH is not None
 
+BASE_MODEL_REPO_ID = os.environ.get(
+    "BASE_MODEL_REPO_ID", "p1atdev/animagine-xl-4.0-bnb-nf4"
+)
+BASE_MODEL_PATH = os.environ.get(
+    "BASE_MODEL_PATH", "animagine-xl-4.0-opt.bnb_nf4.safetensors"
+)
 
-# @spaces.GPU #[uncomment to use ZeroGPU]
-def infer(
-    prompt,
-    negative_prompt,
-    seed,
-    randomize_seed,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    progress=gr.Progress(track_tqdm=True),
-):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    generator = torch.Generator().manual_seed(seed)
+INITIAL_BATCH_SIZE = int(os.environ.get("INITIAL_BATCH_SIZE", 1))
 
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
 
-    return image, seed
+adapter_model_path = get_file_path(ADAPTER_REPO_ID, ADAPTER_MODEL_PATH)
+adapter_config_path = get_file_path(ADAPTER_REPO_ID, ADAPTER_CONFIG_PATH)
 
+base_model_path = get_file_path(BASE_MODEL_REPO_ID, BASE_MODEL_PATH)
 
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-]
+model = load_ip_adapter_model(
+    model_path=base_model_path,
+    config_path=adapter_config_path,
+    adapter_path=adapter_model_path,
+)
+model.to("cuda:0")
 
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 640px;
-}
-"""
 
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image Gradio Template")
+@spaces.GPU
+def on_generate(
+    prompt: str,
+    negative_prompt: str,
+    image: Image.Image | None,
+    width: int,
+    height: int,
+    steps: int,
+    cfg_scale: float,
+    seed: int,
+    randomize_seed: bool = True,
+    num_images: int = 4,
+):
+    if image is not None:
+        image = image.convert("RGB")
 
+    if randomize_seed:
+        seed = random.randint(0, 2147483647)
+
+    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+        images = model.generate(
+            prompt=[prompt] * num_images,  # batch size 4
+            negative_prompt=negative_prompt,
+            reference_image=image,
+            num_inference_steps=steps,
+            cfg_scale=cfg_scale,
+            width=width,
+            height=height,
+            seed=seed,
+            do_offloading=False,
+            device="cuda:0",
+            max_token_length=225,
+            execution_dtype=torch.bfloat16,
+        )
+
+    torch.cuda.empty_cache()
+
+    return images, seed
+
+
+def main():
+    with gr.Blocks() as demo:
         with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
-
-            run_button = gr.Button("Run", scale=0, variant="primary")
-
-        result = gr.Image(label="Result", show_label=False)
-
-        with gr.Accordion("Advanced Settings", open=False):
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
-            )
-
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
-            )
-
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-            with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
+            with gr.Column():
+                prompt = gr.TextArea(
+                    label="Prompt",
+                    value="masterpiece, best quality",
+                    placeholder="masterpiece, best quality",
+                    interactive=True,
                 )
-
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
+                input_image = gr.Image(
+                    label="Reference Image",
+                    type="pil",
+                    height=600,
                 )
 
-            with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=0.0,  # Replace with defaults that work for your model
+                with gr.Accordion("Negative Prompt", open=False):
+                    negative_prompt = gr.TextArea(
+                        label="Negative Prompt",
+                        show_label=False,
+                        value="lowres, bad anatomy, bad hands, text, error, missing finger, extra digits, fewer digits, cropped, worst quality, low quality, low score, bad score, average score, signature, watermark, username, blurry",
+                        interactive=True,
+                    )
+
+                with gr.Row():
+                    width = gr.Slider(
+                        label="Width",
+                        minimum=256,
+                        maximum=2048,
+                        step=128,
+                        value=896,
+                        interactive=True,
+                    )
+                    height = gr.Slider(
+                        label="Height",
+                        minimum=256,
+                        maximum=2048,
+                        step=128,
+                        value=1152,
+                        interactive=True,
+                    )
+
+                with gr.Accordion("Advanced options", open=False):
+                    num_images = gr.Slider(
+                        label="Number of images to generate",
+                        minimum=1,
+                        maximum=8,
+                        step=1,
+                        value=INITIAL_BATCH_SIZE,
+                        interactive=True,
+                    )
+
+                    with gr.Row():
+                        seed = gr.Slider(
+                            label="Seed",
+                            minimum=0,
+                            maximum=2147483647,
+                            step=1,
+                            value=0,
+                        )
+                        randomize_seed = gr.Checkbox(
+                            label="Randomize seed",
+                            value=True,
+                            interactive=True,
+                            scale=1,
+                        )
+
+                    steps = gr.Slider(
+                        label="Inference steps",
+                        minimum=10,
+                        maximum=50,
+                        step=1,
+                        value=25,
+                        interactive=True,
+                    )
+
+                    cfg_scale = gr.Slider(
+                        label="CFG scale",
+                        minimum=3.0,
+                        maximum=8.0,
+                        step=0.5,
+                        value=5.0,
+                        interactive=True,
+                    )
+
+            with gr.Column():
+                generate_button = gr.Button(
+                    "Generate",
+                    variant="primary",
                 )
-
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=2,  # Replace with defaults that work for your model
+                output_image = gr.Gallery(
+                    label="Generated images",
+                    type="pil",
+                    rows=2,
+                    height="768px",
+                    preview=True,
+                    show_label=True,
                 )
 
-        gr.Examples(examples=examples, inputs=[prompt])
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
-        inputs=[
-            prompt,
-            negative_prompt,
-            seed,
-            randomize_seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-        ],
-        outputs=[result, seed],
-    )
+        comment = gr.Markdown(
+            label="Comment",
+            visible=False,
+        )
+
+        gr.Examples(
+            examples=EXAMPLES,
+            inputs=[input_image, prompt, width, height, comment],
+            cache_examples=False,
+        )
+
+        gr.on(
+            triggers=[generate_button.click],
+            fn=on_generate,
+            inputs=[
+                prompt,
+                negative_prompt,
+                input_image,
+                width,
+                height,
+                steps,
+                cfg_scale,
+                seed,
+                randomize_seed,
+                num_images,
+            ],
+            outputs=[output_image, seed],
+        )
 
-if __name__ == "__main__":
     demo.launch()
+
+
+if __name__ == "__main__":
+    main()

example.py

@@ -0,0 +1,31 @@
+EXAMPLES = [
+    # image_path, prompt, width, height, comment
+    [
+        "./examples/marin.png",
+        "masterpiece, best quality",
+        896,
+        1152,
+        "Source: https://ninkoro.jp/chara/marin.html",
+    ],
+    [
+        "./examples/maid.jpg",
+        "1girl, solo, cowboy shot, hands up, heart hands, masterpiece, best quality",
+        896,
+        1152,
+        "Generated with AnimagineXL 4.0",
+    ],
+    [
+        "./examples/nai.png",
+        "masterpiece, best quality",
+        896,
+        1152,
+        "Generated with NovelAI Diffusion v4.5 Full",
+    ],
+    [
+        "./examples/ururu.png",
+        "masterpiece, best quality",
+        896,
+        1152,
+        "Source: https://www.irasutoya.com/2020/06/blog-post_169.html",
+    ],
+]

pyproject.toml

@@ -0,0 +1,24 @@
+[project]
+name = "animaginexl4-0-image-prompt"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "diffusers>=0.35.1",
+    "huggingface-hub>=0.34.4",
+    "safetensors>=0.6.2",
+    "spaces>=0.40.1",
+    "vision-ft",
+]
+
+[dependency-groups]
+dev = [
+    "gradio>=5.43.1",
+ "python-dotenv>=1.1.1",
+ "ruff>=0.12.10",
+ "ty>=0.0.1a19",
+]
+
+[tool.uv.sources]
+vision-ft = { path = "../vision-ft", editable = true }

requirements.txt

@@ -1,6 +1,6 @@
+git+https://github.com/p1atdev/vision-ft@1175478
 accelerate
 diffusers
-invisible_watermark
-torch
 transformers
-xformers
+spaces>=0.40.1
+python-dotenv