notebook fixes

app.ipynb

@@ -17,8 +17,8 @@
    "source": [
     "#| export\n",
     "import gradio as gr\n",
-    "import cf_guidance\n",
-    "import min_diffusion\n",
+    "from cf_guidance import schedules, transforms\n",
+    "from min_diffusion.core import MinimalDiffusion\n",
     "import torch\n",
     "import nbdev"
    ]
@@ -52,7 +52,7 @@
     "}\n",
     "\n",
     "def load_model():\n",
-    "    pipeline = min_diffusion.core.MinimalDiffusion(\n",
+    "    pipeline = MinimalDiffusion(\n",
     "        model_name,\n",
     "        device,\n",
     "        dtype,\n",
@@ -137,7 +137,7 @@
     "    cos_params.update(new_params)\n",
     "    \n",
     "    # return the new cosine schedule\n",
-    "    sched = cf_guidance.schedules.get_cos_sched(**cos_params)\n",
+    "    sched = schedules.get_cos_sched(**cos_params)\n",
     "    return sched\n",
     "\n",
     "\n",
@@ -149,7 +149,7 @@
     "inv_k_sched = [max_val - g + min_val for g in k_sched]\n",
     "\n",
     "# group the schedules \n",
-    "schedules = {\n",
+    "scheds = {\n",
     "    'cosine': {'g': inv_k_sched},\n",
     "    'static': {'g': static_sched},\n",
     "}\n",
@@ -188,9 +188,9 @@
     "    res = []\n",
     "\n",
     "    # generate images with static and dynamic schedules\n",
-    "    for (name,sched) in schedules.items():\n",
+    "    for (name,sched) in scheds.items():\n",
     "        # make the guidance norm\n",
-    "        gtfm = cf_guidance.transforms.GuidanceTfm(sched)\n",
+    "        gtfm = transforms.GuidanceTfm(sched)\n",
     "        # generate the image\n",
     "        with torch.autocast(device), torch.no_grad():\n",
     "            img = pipeline.generate(prompt, gtfm, **generation_kwargs)\n",

app.py

@@ -4,13 +4,13 @@
 __all__ = ['model_name', 'revision', 'dtype', 'device', 'better_vae', 'unet_attn_slice', 'sampler_kls', 'hf_sampler',
            'model_kwargs', 'num_steps', 'height', 'width', 'k_sampler', 'use_karras_sigmas', 'NEG_PROMPT',
            'generation_kwargs', 'baseline_g', 'max_val', 'min_val', 'num_warmup_steps', 'warmup_init_val', 'num_cycles',
-           'k_decay', 'DEFAULT_COS_PARAMS', 'static_sched', 'k_sched', 'inv_k_sched', 'schedules', 'iface',
-           'load_model', 'cos_harness', 'compare_dynamic_guidance']
+           'k_decay', 'DEFAULT_COS_PARAMS', 'static_sched', 'k_sched', 'inv_k_sched', 'scheds', 'iface', 'load_model',
+           'cos_harness', 'compare_dynamic_guidance']
 
 # %% app.ipynb 1
 import gradio as gr
-import cf_guidance
-import min_diffusion
+from cf_guidance import schedules, transforms
+from min_diffusion.core import MinimalDiffusion
 import torch
 import nbdev
 
@@ -36,7 +36,7 @@ model_kwargs = {
 }
 
 def load_model():
-    pipeline = min_diffusion.core.MinimalDiffusion(
+    pipeline = MinimalDiffusion(
         model_name,
         device,
         dtype,
@@ -105,7 +105,7 @@ def cos_harness(new_params: dict) -> dict:
     cos_params.update(new_params)
     
     # return the new cosine schedule
-    sched = cf_guidance.schedules.get_cos_sched(**cos_params)
+    sched = schedules.get_cos_sched(**cos_params)
     return sched
 
 
@@ -117,7 +117,7 @@ k_sched = cos_harness({'k_decay': 0.2})
 inv_k_sched = [max_val - g + min_val for g in k_sched]
 
 # group the schedules 
-schedules = {
+scheds = {
     'cosine': {'g': inv_k_sched},
     'static': {'g': static_sched},
 }
@@ -148,9 +148,9 @@ def compare_dynamic_guidance(prompt):
     res = []
 
     # generate images with static and dynamic schedules
-    for (name,sched) in schedules.items():
+    for (name,sched) in scheds.items():
         # make the guidance norm
-        gtfm = cf_guidance.transforms.GuidanceTfm(sched)
+        gtfm = transforms.GuidanceTfm(sched)
         # generate the image
         with torch.autocast(device), torch.no_grad():
             img = pipeline.generate(prompt, gtfm, **generation_kwargs)