Update app.py

app.py

@@ -6,8 +6,13 @@ import numpy as np
 from diffusers import DiffusionPipeline, DPMSolverSinglestepScheduler, ControlNetModel, StableDiffusionXLControlNetPipeline, AutoencoderKL
 import requests 
 import cv2 
+from uuid import uuid4
+import numpy as np 
+from PIL import Image
 
-pipe = DiffusionPipeline.from_pretrained("mann-e/Mann-E_Dreams", torch_dtype=torch.float16).to("cuda")
+controlnet = ControlNetModel.from_pretrained("diffusers/controlnet-canny-sdxl-1.0", torch_dtype=torch.float16).to("cuda")
+
+pipe = StableDiffusionXLControlNetPipeline.from_pretrained("mann-e/Mann-E_Dreams", controlnet=controlnet, torch_dtype=torch.float16).to("cuda")
 pipe.scheduler = DPMSolverSinglestepScheduler.from_config(pipe.scheduler.config, use_karras_sigmas=True)
 
 #pipe.enable_xformers_memory_efficient_attention()
@@ -16,17 +21,37 @@ pipe.enable_vae_slicing()
 torch.cuda.empty_cache()
 
 @spaces.GPU
-def genie (prompt, negative_prompt, width, height, steps, seed):
+def genie (input_image, prompt, negative_prompt, width, height, steps, seed):
+    #processing the input image
+    res = requests.get(input_image)
+    image_name = f'tmp_{uuid4()}.png'
+
+    if res.ok:
+        with open(image_name, 'wb') as f:
+            f.write(res.content)
+
+    # Canny Edge Detection 
+    image = cv2.imread(image_name)
+    image = np.array(image)
+    image = cv2.Canny(image, 100, 200)
+    image = image[:, :, None]
+    image = np.concatenate([image, image, image], axis=2)
+    #cv2.imwrite('canny.png', image)
+    image = Image.fromarray(image)
+    
+    #generating a new image
     generator = np.random.seed(0) if seed == 0 else torch.manual_seed(seed)
-    int_image = pipe(prompt=prompt, negative_prompt=negative_prompt, width=width, height=height, generator=generator, num_inference_steps=steps, guidance_scale=3.0).images[0]
+    int_image = pipe(prompt=prompt, negative_prompt=negative_prompt, width=width, height=height, generator=generator, num_inference_steps=steps, guidance_scale=3.0, image=image, controlnet_conditioning_scale=conditioning_scale).images[0]
     return int_image
     
-gr.Interface(fn=genie, inputs=[gr.Textbox(label='What you want the AI to generate. 75 Token Limit.'),
+gr.Interface(fn=genie, inputs=[gr.Textbox(label='Base Image URL'),
+    gr.Textbox(label='What you want the AI to generate. 75 Token Limit.'),
     gr.Textbox(label='What you DO NOT want the AI to generate. 75 Token Limit.'),
     gr.Slider(576, maximum=1280, value=768, step=16, label='Width (can go up to 1280, but for square images maximum is 1024x1024)'), 
     gr.Slider(576, maximum=1280, value=768, step=16, label='Height (can go up to 1280, but for square images maximum is 1024x1024)'), 
     gr.Slider(1, maximum=8, value=6, step=1, label='Number of Iterations'), 
     gr.Slider(minimum=0, step=1, maximum=999999999999999999, randomize=True, label="Seed"),
+    gr.Slider(minimum=0, step=0.05, maximum=1, label="Conditioning Scale"),
     ],
     outputs='image', 
     title="Mann-E Dreams",