style

app.py

@@ -1,15 +1,150 @@
+import random
+
+import gradio as gr
+import PIL
 import torch
 from diffusers import StableDiffusionInstructPix2PixPipeline
-import gradio as gr 
-import PIL
 
 cartoonization_id = "instruction-tuning-sd/cartoonizer"
 image_proc_id = "instruction-tuning-sd/low-level-img-proc"
 
+title = "Instruction-tuned Stable Diffusion"
+description = "This Space demonstrates the instruction-tuning on Stable Diffusion. To know more, please check out the [corresponding blog post](https://hf.co/blog/instruction-tuning-sd). Some experimentation tips are available from [the original InstructPix2Pix Space](https://huggingface.co/spaces/timbrooks/instruct-pix2pix)."
+
+
 def load_pipeline(id: str):
-    pipeline = StableDiffusionInstructPix2PixPipeline.from_pretrained(id, torch_dtype=torch.float16).to("cuda")
-    return pipeline 
+    pipeline = StableDiffusionInstructPix2PixPipeline.from_pretrained(
+        id, torch_dtype=torch.float16
+    ).to("cuda")
+    return pipeline
+
+
+def infer_cartoonization(
+    prompt: str,
+    negative_prompt: str,
+    image: PIL.Image.Image,
+    steps: int,
+    img_cfg: float,
+    text_cfg: float,
+    seed: int,
+):
+    pipeline = load_pipeline(cartoonization_id)
+    images = pipeline(
+        prompt,
+        image,
+        negative_prompt=negative_prompt,
+        num_inference_steps=int(steps),
+        image_guidance_scale=img_cfg,
+        guidance_scale=text_cfg,
+        generator=torch.manual_seed(int(seed)),
+    )
+    return images
+
+
+def infer_img_proc(
+    prompt: str,
+    negative_prompt: str,
+    image: PIL.Image.Image,
+    steps: int,
+    img_cfg: float,
+    text_cfg: float,
+    seed: int,
+):
+    pipeline = load_pipeline(image_proc_id)
+    images = pipeline(
+        prompt,
+        image,
+        negative_prompt=negative_prompt,
+        num_inference_steps=int(steps),
+        image_guidance_scale=img_cfg,
+        guidance_scale=text_cfg,
+        generator=torch.manual_seed(int(seed)),
+    )
+    return images
+
+
+examples = [
+    [
+        cartoonization_id,
+        "cartoonize this image",
+        "low quality",
+        "examples/mountain.png",
+        20,
+        1.5,
+        7.5,
+        random.randint(0, 100000),
+    ],
+    [
+        image_proc_id,
+        "derain this image",
+        "low quality",
+        "examples/duck.png",
+        20,
+        1.5,
+        7.5,
+        random.randint(0, 100000),
+    ],
+]
+
+with gr.Blocks(theme="gradio/soft") as demo:
+    gr.Markdown(f"## {title}")
+    gr.Markdown(description)
+
+    with gr.Tab("Cartoonization"):
+        prompt = gr.Textbox(label="Prompt")
+        neg_prompt = gr.Textbox(label="Negative Prompt")
+        input_image = gr.Image(label="Input Image")
+        steps = gr.Slider(minimum=5, maximum=100, step=1)
+        img_cfg = gr.Number(value=1.5, label=f"Image CFG", interactive=True)
+        text_cfg = gr.Number(value=7.5, label=f"Text CFG", interactive=True)
+        seed = gr.Slider(minimum=0, maximum=100000, step=1)
+
+        car_output_gallery = gr.Gallery(
+            label="Generated images", show_label=False, elem_id="gallery"
+        ).style(columns=[2], rows=[2], object_fit="contain", height="auto")
+        submit_btn = gr.Button(value="Submit")
+        all_car_inputs = [prompt, neg_prompt, input_image, img_cfg, text_cfg, seed]
+        submit_btn.click(
+            fn=infer_cartoonization,
+            inputs=all_car_inputs,
+            outputs=[car_output_gallery],
+        )
+
+    with gr.Tab("Low-level image processing"):
+        rompt = gr.Textbox(label="Prompt")
+        neg_prompt = gr.Textbox(label="Negative Prompt")
+        input_image = gr.Image(label="Input Image")
+        steps = gr.Slider(minimum=5, maximum=100, step=1)
+        img_cfg = gr.Number(value=1.5, label=f"Image CFG", interactive=True)
+        text_cfg = gr.Number(value=7.5, label=f"Text CFG", interactive=True)
+        seed = gr.Slider(minimum=0, maximum=100000, step=1)
+
+        img_proc_output_gallery = gr.Gallery(
+            label="Generated images", show_label=False, elem_id="gallery"
+        ).style(columns=[2], rows=[2], object_fit="contain", height="auto")
+        submit_btn = gr.Button(value="Submit")
+        all_img_proc_inputs = [prompt, neg_prompt, input_image, img_cfg, text_cfg, seed]
+        submit_btn.click(
+            fn=infer_img_proc,
+            inputs=all_img_proc_inputs,
+            outputs=[img_proc_output_gallery],
+        )
 
+    gr.Markdown("### Cartoonization example")
+    gr.Examples(
+        [examples[0]],
+        inputs=all_car_inputs,
+        outputs=car_output_gallery,
+        fn=infer_cartoonization,
+        cache_examples=True,
+    )
+    gr.Markdown("### Low-level image processing example")
+    gr.Examples(
+        [examples[0]],
+        inputs=all_img_proc_inputs,
+        outputs=img_proc_output_gallery,
+        fn=infer_img_proc,
+        cache_examples=True,
+    )
 
-def infer(prompt: str, image: PIL.Image.Image, num_inference_steps:int, img_guidance_scale: float):
-    pass
+demo.launch()