{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "initial_id",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-09T09:44:30.641366Z",
     "start_time": "2024-12-09T09:44:11.789050Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import gradio as gr\n",
    "from diffusers import DiffusionPipeline\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from PIL import Image\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ddf33e0d3abacc2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "#append current path\n",
    "sys.path.extend(\"/afs/csail.mit.edu/u/h/huiren/code/diffusion/stable_diffusion/release/hf_demo\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "643e49fd601daf8f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-09T09:44:35.790962Z",
     "start_time": "2024-12-09T09:44:35.779496Z"
    }
   },
   "outputs": [],
   "source": [
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\"\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e03aae2a4e5676dd",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-09T09:44:44.157412Z",
     "start_time": "2024-12-09T09:44:37.138452Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/data/vision/torralba/selfmanaged/torralba/scratch/jomat/sam_dataset/miniforge3/envs/diffusion/lib/python3.9/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "acc42f294243439798e4d77d1a59296d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading pipeline components...:   0%|          | 0/7 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pipe = DiffusionPipeline.from_pretrained(\"rhfeiyang/art-free-diffusion-v1\",\n",
    "                                         torch_dtype=dtype).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "83916bc68ff5d914",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-09T09:44:52.694399Z",
     "start_time": "2024-12-09T09:44:44.210695Z"
    }
   },
   "outputs": [],
   "source": [
    "from inference import get_lora_network, inference, get_validation_dataloader\n",
    "lora_map = {\n",
    "    \"None\": \"None\",\n",
    "    \"Andre Derain (fauvism)\": \"andre-derain_subset1\",\n",
    "    \"Vincent van Gogh (post impressionism)\": \"van_gogh_subset1\",\n",
    "    \"Andy Warhol (pop art)\": \"andy_subset1\",\n",
    "    \"Walter Battiss\": \"walter-battiss_subset2\",\n",
    "    \"Camille Corot (realism)\": \"camille-corot_subset1\",\n",
    "    \"Claude Monet (impressionism)\": \"monet_subset2\",\n",
    "    \"Pablo Picasso (cubism)\": \"picasso_subset1\",\n",
    "    \"Jackson Pollock\": \"jackson-pollock_subset1\",\n",
    "    \"Gerhard Richter (abstract expressionism)\": \"gerhard-richter_subset1\",\n",
    "    \"M.C. Escher\": \"m.c.-escher_subset1\",\n",
    "    \"Albert Gleizes\": \"albert-gleizes_subset1\",\n",
    "    \"Hokusai (ukiyo-e)\": \"katsushika-hokusai_subset1\",\n",
    "    \"Wassily Kandinsky\": \"kandinsky_subset1\",\n",
    "    \"Gustav Klimt (art nouveau)\": \"klimt_subset3\",\n",
    "    \"Roy Lichtenstein\": \"roy-lichtenstein_subset1\",\n",
    "    \"Henri Matisse (abstract expressionism)\": \"henri-matisse_subset1\",\n",
    "    \"Joan Miro\": \"joan-miro_subset2\",\n",
    "}\n",
    "\n",
    "\n",
    "\n",
    "def demo_inference_gen_artistic(adapter_choice:str, prompt:str, seed:int=0, steps=50, guidance_scale=7.5, adapter_scale=1.0):\n",
    "    adapter_path = lora_map[adapter_choice]\n",
    "    if adapter_path not in [None, \"None\"]:\n",
    "        adapter_path = f\"data/Art_adapters/{adapter_path}/adapter_alpha1.0_rank1_all_up_1000steps.pt\"\n",
    "        style_prompt=\"sks art\"\n",
    "    else:\n",
    "        style_prompt=None\n",
    "    prompts = [prompt]\n",
    "    infer_loader = get_validation_dataloader(prompts,num_workers=0)\n",
    "    network = get_lora_network(pipe.unet, adapter_path, weight_dtype=dtype)[\"network\"]\n",
    "\n",
    "    pred_images = inference(network, pipe.tokenizer, pipe.text_encoder, pipe.vae, pipe.unet, pipe.scheduler, infer_loader,\n",
    "                            height=512, width=512, scales=[adapter_scale],\n",
    "                            save_dir=None, seed=seed,steps=steps, guidance_scale=guidance_scale,\n",
    "                            start_noise=-1, show=False, style_prompt=style_prompt, no_load=True,\n",
    "                            from_scratch=True, device=device, weight_dtype=dtype)[0][1.0][0]\n",
    "    return pred_images\n",
    "\n",
    "\n",
    "def demo_inference_gen_ori( prompt:str, seed:int=0, steps=50, guidance_scale=7.5):\n",
    "    style_prompt=None\n",
    "    prompts = [prompt]\n",
    "    infer_loader = get_validation_dataloader(prompts,num_workers=0)\n",
    "    network = get_lora_network(pipe.unet, \"None\", weight_dtype=dtype)[\"network\"]\n",
    "\n",
    "    pred_images = inference(network, pipe.tokenizer, pipe.text_encoder, pipe.vae, pipe.unet, pipe.scheduler, infer_loader,\n",
    "                            height=512, width=512, scales=[0.0],\n",
    "                            save_dir=None, seed=seed,steps=steps, guidance_scale=guidance_scale,\n",
    "                            start_noise=-1, show=False, style_prompt=style_prompt, no_load=True,\n",
    "                            from_scratch=True, device=device, weight_dtype=dtype)[0][0.0][0]\n",
    "    return pred_images\n",
    "\n",
    "\n",
    "\n",
    "def demo_inference_stylization_ori(ref_image, prompt:str, seed:int=0, steps=50, guidance_scale=7.5, start_noise=800):\n",
    "    style_prompt=None\n",
    "    prompts = [prompt]\n",
    "    # convert np to pil\n",
    "    ref_image = [Image.fromarray(ref_image)]\n",
    "    network = get_lora_network(pipe.unet, \"None\", weight_dtype=dtype)[\"network\"]\n",
    "    infer_loader = get_validation_dataloader(prompts, ref_image,num_workers=0)\n",
    "    pred_images = inference(network, pipe.tokenizer, pipe.text_encoder, pipe.vae, pipe.unet, pipe.scheduler, infer_loader,\n",
    "                            height=512, width=512, scales=[0.0],\n",
    "                            save_dir=None, seed=seed,steps=steps, guidance_scale=guidance_scale,\n",
    "                            start_noise=start_noise, show=False, style_prompt=style_prompt, no_load=True,\n",
    "                            from_scratch=False, device=device, weight_dtype=dtype)[0][0.0][0]\n",
    "    return pred_images\n",
    "\n",
    "\n",
    "def demo_inference_stylization_artistic(ref_image, adapter_choice:str, prompt:str, seed:int=0, steps=50, guidance_scale=7.5, adapter_scale=1.0,start_noise=800):\n",
    "    adapter_path = lora_map[adapter_choice]\n",
    "    if adapter_path not in [None, \"None\"]:\n",
    "        adapter_path = f\"data/Art_adapters/{adapter_path}/adapter_alpha1.0_rank1_all_up_1000steps.pt\"\n",
    "        style_prompt=\"sks art\"\n",
    "    else:\n",
    "        style_prompt=None\n",
    "    prompts = [prompt]\n",
    "    # convert np to pil\n",
    "    ref_image = [Image.fromarray(ref_image)]\n",
    "    network = get_lora_network(pipe.unet, adapter_path, weight_dtype=dtype)[\"network\"]\n",
    "    infer_loader = get_validation_dataloader(prompts, ref_image,num_workers=0)\n",
    "    pred_images = inference(network, pipe.tokenizer, pipe.text_encoder, pipe.vae, pipe.unet, pipe.scheduler, infer_loader,\n",
    "                            height=512, width=512, scales=[adapter_scale],\n",
    "                            save_dir=None, seed=seed,steps=steps, guidance_scale=guidance_scale,\n",
    "                            start_noise=start_noise, show=False, style_prompt=style_prompt, no_load=True,\n",
    "                            from_scratch=False, device=device, weight_dtype=dtype)[0][1.0][0]\n",
    "    return pred_images\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "aa33e9d104023847",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-10T02:56:13.419303Z",
     "start_time": "2024-12-10T02:56:13.002796Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running on local URL:  http://127.0.0.1:7869\n",
      "\n",
      "Thanks for being a Gradio user! If you have questions or feedback, please join our Discord server and chat with us: https://discord.gg/feTf9x3ZSB\n",
      "Running on public URL: https://0fd0c028b349b76a72.gradio.live\n",
      "\n",
      "This share link expires in 72 hours. For free permanent hosting and GPU upgrades, run `gradio deploy` from Terminal to deploy to Spaces (https://huggingface.co/spaces)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><iframe src=\"https://0fd0c028b349b76a72.gradio.live\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": []
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "block = gr.Blocks()\n",
    "# Direct infer\n",
    "# Direct infer\n",
    "with block:\n",
    "    with gr.Group():\n",
    "        gr.Markdown(\" # Art-Free Diffusion Demo\")\n",
    "        gr.Markdown(\"(More features in development...)\")\n",
    "        with gr.Row():\n",
    "            text = gr.Textbox(\n",
    "                label=\"Enter your prompt(long and detailed would be better):\",\n",
    "                max_lines=2,\n",
    "                placeholder=\"Enter your prompt(long and detailed would be better)\",\n",
    "                container=True,\n",
    "                value=\"Park with cherry blossom trees, picnicker’s and a clear blue pond.\",\n",
    "            )\n",
    "\n",
    "        with gr.Tab('Generation'):\n",
    "            with gr.Row():\n",
    "                with gr.Column():\n",
    "                    # gr.Markdown(\"## Art-Free Generation\")\n",
    "                    # gr.Markdown(\"Generate images from text prompts.\")\n",
    "\n",
    "                    gallery_gen_ori = gr.Image(\n",
    "                        label=\"W/O Adapter\",\n",
    "                        show_label=True,\n",
    "                        elem_id=\"gallery\",\n",
    "                        height=\"auto\"\n",
    "                    )\n",
    "\n",
    "\n",
    "                with gr.Column():\n",
    "                    # gr.Markdown(\"## Art-Free Generation\")\n",
    "                    # gr.Markdown(\"Generate images from text prompts.\")\n",
    "                    gallery_gen_art = gr.Image(\n",
    "                        label=\"W/ Adapter\",\n",
    "                        show_label=True,\n",
    "                        elem_id=\"gallery\",\n",
    "                        height=\"auto\"\n",
    "                    )\n",
    "\n",
    "\n",
    "            with gr.Row():\n",
    "                btn_gen_ori = gr.Button(\"Art-Free Generate\", scale=1)\n",
    "                btn_gen_art = gr.Button(\"Artistic Generate\", scale=1)\n",
    "\n",
    "\n",
    "        with gr.Tab('Stylization'):\n",
    "            with gr.Row():\n",
    "\n",
    "                with gr.Column():\n",
    "                    # gr.Markdown(\"## Art-Free Generation\")\n",
    "                    # gr.Markdown(\"Generate images from text prompts.\")\n",
    "\n",
    "                    gallery_stylization_ref = gr.Image(\n",
    "                        label=\"Ref Image\",\n",
    "                        show_label=True,\n",
    "                        elem_id=\"gallery\",\n",
    "                        height=\"auto\",\n",
    "                        scale=1,\n",
    "                    )\n",
    "                with gr.Column(scale=2):\n",
    "                    with gr.Row():\n",
    "                        with gr.Column():\n",
    "                            # gr.Markdown(\"## Art-Free Generation\")\n",
    "                            # gr.Markdown(\"Generate images from text prompts.\")\n",
    "    \n",
    "                            gallery_stylization_ori = gr.Image(\n",
    "                                label=\"W/O Adapter\",\n",
    "                                show_label=True,\n",
    "                                elem_id=\"gallery\",\n",
    "                                height=\"auto\",\n",
    "                                scale=1,\n",
    "                            )\n",
    "    \n",
    "    \n",
    "                        with gr.Column():\n",
    "                            # gr.Markdown(\"## Art-Free Generation\")\n",
    "                            # gr.Markdown(\"Generate images from text prompts.\")\n",
    "                            gallery_stylization_art = gr.Image(\n",
    "                                label=\"W/ Adapter\",\n",
    "                                show_label=True,\n",
    "                                elem_id=\"gallery\",\n",
    "                                height=\"auto\",\n",
    "                                scale=1,\n",
    "                            )\n",
    "                    start_timestep = gr.Slider(label=\"Adapter Timestep\", minimum=0, maximum=1000, value=800, step=1)\n",
    "            with gr.Row():\n",
    "                btn_style_ori = gr.Button(\"Art-Free Stylization\", scale=1)\n",
    "                btn_style_art = gr.Button(\"Artistic Stylization\", scale=1)\n",
    "\n",
    "\n",
    "        with gr.Row():\n",
    "            # with gr.Column():\n",
    "            # samples = gr.Slider(label=\"Images\", minimum=1, maximum=4, value=1, step=1, scale=1)\n",
    "            scale = gr.Slider(\n",
    "                label=\"Guidance Scale\", minimum=0, maximum=20, value=7.5, step=0.1\n",
    "            )\n",
    "            # with gr.Column():\n",
    "            adapter_choice = gr.Dropdown(\n",
    "                label=\"Select Art Adapter\",\n",
    "                choices=[ \"Andre Derain (fauvism)\",\"Vincent van Gogh (post impressionism)\",\"Andy Warhol (pop art)\",\n",
    "                          \"Camille Corot (realism)\", \"Claude Monet (impressionism)\", \"Pablo Picasso (cubism)\", \"Gerhard Richter (abstract expressionism)\",\n",
    "                          \"Hokusai (ukiyo-e)\", \"Gustav Klimt (art nouveau)\", \"Henri Matisse (abstract expressionism)\",\n",
    "                          \"Walter Battiss\", \"Jackson Pollock\",  \"M.C. Escher\", \"Albert Gleizes\",  \"Wassily Kandinsky\",\n",
    "                          \"Roy Lichtenstein\", \"Joan Miro\"\n",
    "                          ],\n",
    "                value=\"Andre Derain (fauvism)\",\n",
    "                scale=1\n",
    "            )\n",
    "\n",
    "        with gr.Row():\n",
    "            steps = gr.Slider(label=\"Steps\", minimum=1, maximum=50, value=20, step=1)\n",
    "            adapter_scale = gr.Slider(label=\"Stylization Scale\", minimum=0, maximum=1.5, value=1., step=0.1, scale=1)\n",
    "\n",
    "        with gr.Row():\n",
    "            seed = gr.Slider(label=\"Seed\",minimum=0,maximum=2147483647,step=1,randomize=True,scale=1)\n",
    "\n",
    "\n",
    "        gr.on([btn_gen_ori.click], demo_inference_gen_ori, inputs=[text, seed, steps, scale], outputs=gallery_gen_ori)\n",
    "        gr.on([btn_gen_art.click], demo_inference_gen_artistic, inputs=[adapter_choice, text, seed, steps, scale, adapter_scale], outputs=gallery_gen_art)\n",
    "\n",
    "        gr.on([btn_style_ori.click], demo_inference_stylization_ori, inputs=[gallery_stylization_ref, text, seed, steps, scale, start_timestep], outputs=gallery_stylization_ori)\n",
    "        gr.on([btn_style_art.click], demo_inference_stylization_artistic, inputs=[gallery_stylization_ref, adapter_choice, text, seed, steps, scale, adapter_scale, start_timestep], outputs=gallery_stylization_art)\n",
    "\n",
    "block.launch(share=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3239c12167a5f2cd",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}