Update app.py

app.py

@@ -14,9 +14,7 @@ import torch as th
 from composable_diffusion.download import download_model
 from composable_diffusion.model_creation import create_model_and_diffusion as create_model_and_diffusion_for_clevr
 from composable_diffusion.model_creation import model_and_diffusion_defaults as model_and_diffusion_defaults_for_clevr
-
-from torch import autocast
-from composable_stable_diffusion_pipeline import ComposableStableDiffusionPipeline
+from composable_diffusion.composable_stable_diffusion.pipeline_composable_stable_diffusion import ComposableStableDiffusionPipeline
 
 # This notebook supports both CPU and GPU.
 # On CPU, generating one sample may take on the order of 20 minutes.
@@ -24,7 +22,6 @@ from composable_stable_diffusion_pipeline import ComposableStableDiffusionPipeli
 
 has_cuda = th.cuda.is_available()
 device = th.device('cpu' if not th.cuda.is_available() else 'cuda')
-print(device)
 
 # init stable diffusion model
 pipe = ComposableStableDiffusionPipeline.from_pretrained(
@@ -32,9 +29,7 @@ pipe = ComposableStableDiffusionPipeline.from_pretrained(
     use_auth_token=st.secrets["USER_TOKEN"]
 ).to(device)
 
-def dummy(images, **kwargs):
-    return images, False
-pipe.safety_checker = dummy
+pipe.safety_checker = None
 
 # create model for CLEVR Objects
 clevr_options = model_and_diffusion_defaults_for_clevr()
@@ -68,9 +63,14 @@ clevr_model.load_state_dict(th.load(download_model('clevr_pos'), device))
 print('total clevr_pos parameters', sum(x.numel() for x in clevr_model.parameters()))
 
 
-def compose_clevr_objects(prompt, guidance_scale, steps):
-    coordinates = [[float(x.split(',')[0].strip()), float(x.split(',')[1].strip())]
-                   for x in prompt.split('|')]
+def compose_clevr_objects(prompt, weights, steps):
+    weights = [float(x.strip()) for x in weights.split('|')]
+    weights = th.tensor(weights, device=device).reshape(-1, 1, 1, 1)
+    coordinates = [
+        [
+            float(x.split(',')[0].strip()), float(x.split(',')[1].strip())]
+        for x in prompt.split('|')
+    ]
     coordinates += [[-1, -1]]  # add unconditional score label
     batch_size = 1
 
@@ -83,9 +83,9 @@ def compose_clevr_objects(prompt, guidance_scale, steps):
         model_out = clevr_model(combined, ts, **kwargs)
         eps, rest = model_out[:, :3], model_out[:, 3:]
         masks = kwargs.get('masks')
-        cond_eps = eps[masks].mean(dim=0, keepdim=True)
-        uncond_eps = eps[~masks].mean(dim=0, keepdim=True)
-        half_eps = uncond_eps + guidance_scale * (cond_eps - uncond_eps)
+        cond_eps = eps[masks]
+        uncond_eps = eps[~masks]
+        half_eps = uncond_eps + (weights * (cond_eps - uncond_eps)).sum(dim=0, keepdims=True)
         eps = th.cat([half_eps] * x_t.size(0), dim=0)
         return th.cat([eps, rest], dim=1)
 
@@ -116,38 +116,38 @@ def compose_clevr_objects(prompt, guidance_scale, steps):
     return out_img
 
 
-def stable_diffusion_compose(prompt, scale, steps, weights, seed):
+def stable_diffusion_compose(prompt, steps, weights, seed):
     generator = th.Generator("cuda").manual_seed(int(seed))
-    with autocast('cpu' if not th.cuda.is_available() else 'cuda'):
-        image = pipe(prompt, guidance_scale=scale, num_inference_steps=steps,
-                     weights=weights, generator=generator).images[0]
-        image.save(f'{"_".join(prompt.split())}.png')
-        return image
+    image = pipe(prompt, guidance_scale=7.5, num_inference_steps=steps,
+                 weights=weights, generator=generator).images[0]
+    image.save(f'{"_".join(prompt.split())}.png')
+    return image
 
 
-def compose(prompt, weights, version, guidance_scale, steps, seed):
+def compose(prompt, weights, version, steps, seed):
     try:
         with th.no_grad():
             if version == 'Stable_Diffusion_1v_4':
-                return stable_diffusion_compose(prompt, guidance_scale, steps, weights, seed)
+                res = stable_diffusion_compose(prompt, steps, weights, seed)
+                return res
             else:
-                return compose_clevr_objects(prompt, guidance_scale, steps)
+                return compose_clevr_objects(prompt, weights, steps)
     except Exception as e:
         print(e)
         return None
 
 examples_1 = "A castle in a forest | grainy, fog"
-examples_2 = 'A blue sky  | A mountain in the horizon | Cherry Blossoms in front of the mountain'
 examples_3 = '0.1, 0.5 | 0.3, 0.5 | 0.5, 0.5 | 0.7, 0.5 | 0.9, 0.5'
 examples_5 = 'a white church | lightning in the background'
 examples_6 = 'mystical trees | A dark magical pond | dark'
 examples_7 = 'A lake | A mountain  | Cherry Blossoms next to the lake'
 examples = [
-            [examples_1, "1 | -1", 'Stable_Diffusion_1v_4', 15, 50, 0],
-            [examples_7, "1 | 1 | 1", 'Stable_Diffusion_1v_4', 15, 50, 3],
-            [examples_5, "1 | 1", 'Stable_Diffusion_1v_4', 15, 50, 0],
-            [examples_6, "1 | 1 | -1", 'Stable_Diffusion_1v_4', 15, 50, 2],
-            [examples_3, "1 | 1 | 1 | 1 | 1", 'CLEVR Objects', 10, 100, 0]
+            [examples_6, "7.5 | 7.5 | -7.5", 'Stable_Diffusion_1v_4', 50, 8],
+            [examples_6, "7.5 | 7.5 | 7.5", 'Stable_Diffusion_1v_4', 50, 8],
+            [examples_1, "7.5 | -7.5", 'Stable_Diffusion_1v_4', 50, 0],
+            [examples_7, "7.5 | 7.5 | 7.5", 'Stable_Diffusion_1v_4', 50, 3],
+            [examples_5, "7.5 | 7.5", 'Stable_Diffusion_1v_4', 50, 0],
+            [examples_3, "7.5 | 7.5 | 7.5 | 7.5 | 7.5", 'CLEVR Objects', 100, 0]
 ]
 
 title = 'Compositional Visual Generation with Composable Diffusion Models'
@@ -156,13 +156,12 @@ description = '<p>Our conjunction and negation (a.k.a. negative prompts) operato
 iface = gr.Interface(compose,
                      inputs=[
                          gr.Textbox(label='prompt', value='mystical trees | A dark magical pond | dark'),
-                         gr.Textbox(label='weights', value='1 | 1 | -1'),
+                         gr.Textbox(label='weights', value='7.5 | 7.5 | -7.5'),
                          gr.Radio(['Stable_Diffusion_1v_4', 'CLEVR Objects'], type="value", label='version', value='Stable_Diffusion_1v_4'),
-                         gr.Slider(2, 30, value=15),
                          gr.Slider(10, 200, value=50),
                          gr.Number(2)
                      ],
                      outputs='image', cache_examples=False,
                      title=title, description=description, examples=examples)
 
-iface.launch()
+iface.launch()