add init demo

README.md

@@ -1,12 +1,11 @@
 ---
-title: Noisy Style
-emoji: 🏃
-colorFrom: yellow
-colorTo: red
+title: Noisy-Style
+emoji: 🎨
+colorFrom: blue
+colorTo: pink
 sdk: gradio
-sdk_version: 4.44.0
+sdk_version: 4.37.2
 app_file: app.py
 pinned: false
+license: cc-by-nc-4.0
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

main.py

@@ -0,0 +1,372 @@
+import os
+import random
+from typing import Optional
+
+import gradio as gr
+import numpy as np
+import PIL.Image
+import spaces
+import torch
+from diffusers import (AutoencoderKL, DDIMInverseScheduler, DDIMScheduler,
+                       StableDiffusionXLPipeline)
+from torchvision.transforms import ToTensor 
+
+# pyright: reportPrivateImportUsage=false
+
+
+DESCRIPTION = f"""
+    # 🎨 Noisy-Style 🎨  
+    This is an interactive demo of noisy DDIM inversion capabilities on top of Instant-Style styling method 
+    
+    This method proposed in [Controllability of diffusion models]() *by Eyal Benaroche, Clément Chadebec, Onur Tasar, and Benjamin Aubin* from Jasper Research in the context of Eyal's internship with Ecole Polytechnique.
+
+    A style benchmark : [style-bench](https://gojasper.github.io/style-bench) was also provided to facilitate evaluation of diffusion models for styling purposes.
+    """
+
+OPEN_SOURCE_PROMO = f"""
+    If you enjoy the space, please also promote *open-source* by giving a ⭐ to our repo [![GitHub Stars](https://img.shields.io/github/stars/gojasper/style-bench?style=social)](https://github.com/gojasper/style-bench)
+    """
+
+DISCLAIMER = f"""
+    This demo is only for research purpose. Users are solely responsible for any content they create, and it is their obligation to ensure that it adheres to appropriate and ethical standards. """
+
+if not torch.cuda.is_available():
+    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
+
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1024"))
+USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE") == "1"
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+if gr.NO_RELOAD:
+    if torch.cuda.is_available():
+        vae = AutoencoderKL.from_pretrained(
+            "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
+        )
+        pipe = StableDiffusionXLPipeline.from_pretrained(
+            "stabilityai/stable-diffusion-xl-base-1.0",
+            vae=vae,
+            torch_dtype=torch.float16,
+            use_safetensors=True,
+            variant="fp16",
+        )
+        pipe.load_ip_adapter(
+            "h94/IP-Adapter",
+            subfolder="sdxl_models",
+            weight_name="ip-adapter_sdxl.safetensors"
+        )
+        pipe.to(device)
+
+        forward_scheduler = DDIMScheduler.from_pretrained(
+            "stabilityai/stable-diffusion-xl-base-1.0", subfolder="scheduler"
+        )
+        invert_scheduler = DDIMInverseScheduler(**forward_scheduler.config)
+
+css = """
+    h1 {
+        text-align: center;
+        display:block;
+    }
+    p {
+        text-align: justify;
+        display:block;
+    }
+"""
+
+
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
+
+
+def img_to_latents(x: torch.Tensor, vae: AutoencoderKL):
+    x = 2.0 * x - 1.0
+    posterior = vae.encode(x).latent_dist
+    latents = posterior.mean * 0.18215
+    return latents
+
+
+def invert_image(model, image: np.ndarray, n_steps: int, width:int, height:int):
+
+    model.scheduler = invert_scheduler
+
+    image = PIL.Image.fromarray(image).resize((width, height))
+    image_tensor = ToTensor()(image).to(model.device, dtype=torch.float16)
+    image_tensor = image_tensor.unsqueeze(0)
+    latent = img_to_latents(image_tensor, model.vae)
+
+    print(latent)
+    model.set_ip_adapter_scale(0)
+
+    inv_latents = model(
+        prompt="",
+        negative_prompt="",
+        ip_adapter_image=image,
+        guidance_scale=1.0,
+        output_type="latent",
+        return_dict=False,
+        num_inference_steps=n_steps,
+        latents=latent,
+    )[0]
+
+    return inv_latents
+
+
+@spaces.GPU
+def generate(
+    prompt: str,
+    negative_prompt: str = "",
+    prompt_2: str = "",
+    negative_prompt_2: str = "",
+    use_negative_prompt: bool = False,
+    use_prompt_2: bool = False,
+    use_negative_prompt_2: bool = False,
+    seed: int = 0,
+    width: int = 1024,
+    height: int = 1024,
+    guidance_scale_base: float = 5.0,
+    num_inference_steps_base: int = 25,
+    style_image_value = None,
+    noise_scale: float = 1.5,
+) -> PIL.Image.Image:
+    torch.manual_seed(seed)
+
+    if not use_negative_prompt:
+        negative_prompt = None  # type: ignore
+    if not use_prompt_2:
+        prompt_2 = None  # type: ignore
+    if not use_negative_prompt_2:
+        negative_prompt_2 = None  # type: ignore
+
+    # Add scaled noise to the latent
+    noise = torch.randn(1, 4, width // 8, height // 8).to(device, dtype=torch.float16)
+
+    # Invert the image and get the latent
+    if style_image_value is not None:
+        latent = invert_image(pipe, style_image_value, 30, width, height)
+        print("Image was inverted")
+        print(latent)
+
+        latent = latent + noise_scale * noise
+        latent = latent / torch.sqrt(torch.tensor(1 + noise_scale**2).to(device, dtype=torch.float16))
+
+    else:
+        latent = noise
+
+    print("Noise added")
+    print(latent)
+
+    scale = {
+        "up": {"block_0": [0.0, 1.0, 0.0]},
+    }
+    pipe.set_ip_adapter_scale(scale)
+
+    pipe.scheduler = forward_scheduler
+    image = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        ip_adapter_image=style_image_value,
+        latents=latent,
+        prompt_2=prompt_2,
+        negative_prompt_2=negative_prompt_2,
+        guidance_scale=guidance_scale_base,
+        num_inference_steps=num_inference_steps_base,
+        output_type="pil",
+    ).images[0]
+
+    return image
+
+
+examples = [
+    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
+    "An astronaut riding a green horse",
+]
+
+with gr.Blocks(css=css) as demo:
+    gr.Markdown(DESCRIPTION)
+
+    gr.Markdown(OPEN_SOURCE_PROMO)
+
+    with gr.Row():
+
+        with gr.Blocks():
+
+            with gr.Column():
+
+                style_image = gr.Image()
+
+                noise_scale = gr.Slider(
+                    label="Noise Scale",
+                    minimum=0,
+                    maximum=5,
+                    step=0.1,
+                    value=1.5,
+                )
+
+        with gr.Blocks():
+
+            with gr.Column():
+                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)
+
+                result = gr.Image(label="Result", show_label=False)
+
+    with gr.Accordion("Advanced options", open=False):
+        with gr.Row():
+            use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=False)
+            use_prompt_2 = gr.Checkbox(label="Use prompt 2", value=False)
+            use_negative_prompt_2 = gr.Checkbox(
+                label="Use negative prompt 2", value=False
+            )
+        negative_prompt = gr.Text(
+            label="Negative prompt",
+            max_lines=1,
+            placeholder="Enter a negative prompt",
+            visible=False,
+        )
+        prompt_2 = gr.Text(
+            label="Prompt 2",
+            max_lines=1,
+            placeholder="Enter your prompt",
+            visible=False,
+        )
+        negative_prompt_2 = gr.Text(
+            label="Negative prompt 2",
+            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,
+            )
+            height = gr.Slider(
+                label="Height",
+                minimum=256,
+                maximum=MAX_IMAGE_SIZE,
+                step=32,
+                value=1024,
+            )
+
+        with gr.Row():
+            guidance_scale_base = gr.Slider(
+                label="Guidance scale for base",
+                minimum=1,
+                maximum=20,
+                step=0.1,
+                value=5.0,
+            )
+            num_inference_steps_base = gr.Slider(
+                label="Number of inference steps for base",
+                minimum=10,
+                maximum=100,
+                step=1,
+                value=25,
+            )
+        with gr.Row(visible=False) as refiner_params:
+            guidance_scale_refiner = gr.Slider(
+                label="Guidance scale for refiner",
+                minimum=1,
+                maximum=20,
+                step=0.1,
+                value=5.0,
+            )
+            num_inference_steps_refiner = gr.Slider(
+                label="Number of inference steps for refiner",
+                minimum=10,
+                maximum=100,
+                step=1,
+                value=25,
+            )
+
+    gr.Examples(
+        examples=examples,
+        inputs=prompt,
+        outputs=result,
+        fn=generate,
+    )
+
+    gr.Markdown("## Disclaimer")
+    gr.Markdown(DISCLAIMER)
+
+    use_negative_prompt.change(
+        fn=lambda x: gr.update(visible=x),
+        inputs=use_negative_prompt,
+        outputs=negative_prompt,
+        queue=False,
+        api_name=False,
+    )
+    use_prompt_2.change(
+        fn=lambda x: gr.update(visible=x),
+        inputs=use_prompt_2,
+        outputs=prompt_2,
+        queue=False,
+        api_name=False,
+    )
+    use_negative_prompt_2.change(
+        fn=lambda x: gr.update(visible=x),
+        inputs=use_negative_prompt_2,
+        outputs=negative_prompt_2,
+        queue=False,
+        api_name=False,
+    )
+
+    gr.on(
+        triggers=[
+            prompt.submit,
+            negative_prompt.submit,
+            prompt_2.submit,
+            negative_prompt_2.submit,
+            run_button.click,
+        ],
+        fn=randomize_seed_fn,
+        inputs=[seed, randomize_seed],
+        outputs=seed,
+        queue=False,
+        api_name=False,
+    ).then(
+        fn=generate,
+        inputs=[
+            prompt,
+            negative_prompt,
+            prompt_2,
+            negative_prompt_2,
+            use_negative_prompt,
+            use_prompt_2,
+            use_negative_prompt_2,
+            seed,
+            width,
+            height,
+            guidance_scale_base,
+            num_inference_steps_base,
+            style_image,
+            noise_scale,
+        ],
+        outputs=result,
+        api_name="run",
+    )
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,13 @@
+accelerate
+diffusers
+gradio==4.37.2
+numpy==1.26.4
+spaces
+--extra-index-url https://download.pytorch.org/whl/cu118
+torch==2.0.1
+torchvision
+transformers >= 4.34.0
+xformers
+ftfy
+peft==0.6.0
+optimum