Muhammadreza's picture
Update app.py
53d7a5e verified
import spaces
import gradio as gr
import torch
import modin.pandas as pd
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
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()
pipe.enable_vae_slicing()
torch.cuda.empty_cache()
@spaces.GPU
def genie (input_image, prompt, negative_prompt, width, height, steps, seed, conditioning_scale):
#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)
conditioning_scale = float(conditioning_scale)
#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=4, image=image, controlnet_conditioning_scale=conditioning_scale).images[0]
return int_image
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 w/ ControlNet",
description="Mann-E Dreams <br><br><b>WARNING: This model is capable of producing NSFW (Softcore) images.</b><br><br>In case you don't want a base image, just paste link to an image and put conditioning scale to 0",
article = "").launch(debug=True, max_threads=80, show_error=True)