Spaces:
Running
on
A10G
Running
on
A10G
| import functools | |
| import random | |
| import gradio as gr | |
| import torch | |
| from fabric.generator import AttentionBasedGenerator | |
| #model_name = "dreamlike-art/dreamlike-photoreal-2.0" | |
| model_name = "" | |
| model_ckpt = "https://huggingface.co/Lykon/DreamShaper/blob/main/DreamShaper_7_pruned.safetensors" | |
| class GeneratorWrapper: | |
| def __init__(self, model_name=None, model_ckpt=None): | |
| self.model_name = model_name if model_name else None | |
| self.model_ckpt = model_ckpt if model_ckpt else None | |
| self.dtype = torch.float16 if torch.cuda.is_available() else torch.float32 | |
| self.device = "cuda" if torch.cuda.is_available() else "cpu" | |
| self.reload() | |
| def generate(self, *args, **kwargs): | |
| if not hasattr(self, "generator"): | |
| self.reload() | |
| return self.generator.generate(*args, **kwargs) | |
| def to(self, device): | |
| return self.generator.to(device) | |
| def reload(self): | |
| if hasattr(self, "generator"): | |
| del self.generator | |
| if self.device == "cuda": | |
| torch.cuda.empty_cache() | |
| self.generator = AttentionBasedGenerator( | |
| model_name=self.model_name, | |
| model_ckpt=self.model_ckpt, | |
| torch_dtype=self.dtype, | |
| ).to(self.device) | |
| generator = GeneratorWrapper(model_name, model_ckpt) | |
| css = """ | |
| .btn-green { | |
| background-image: linear-gradient(to bottom right, #86efac, #22c55e) !important; | |
| border-color: #22c55e !important; | |
| color: #166534 !important; | |
| } | |
| .btn-green:hover { | |
| background-image: linear-gradient(to bottom right, #86efac, #86efac) !important; | |
| } | |
| .btn-red { | |
| background: linear-gradient(to bottom right, #fda4af, #fb7185) !important; | |
| border-color: #fb7185 !important; | |
| color: #9f1239 !important; | |
| } | |
| .btn-red:hover {background: linear-gradient(to bottom right, #fda4af, #fda4af) !important;} | |
| /*****/ | |
| .dark .btn-green { | |
| background-image: linear-gradient(to bottom right, #047857, #065f46) !important; | |
| border-color: #047857 !important; | |
| color: #ffffff !important; | |
| } | |
| .dark .btn-green:hover { | |
| background-image: linear-gradient(to bottom right, #047857, #047857) !important; | |
| } | |
| .dark .btn-red { | |
| background: linear-gradient(to bottom right, #be123c, #9f1239) !important; | |
| border-color: #be123c !important; | |
| color: #ffffff !important; | |
| } | |
| .dark .btn-red:hover {background: linear-gradient(to bottom right, #be123c, #be123c) !important;} | |
| """ | |
| def generate_fn( | |
| feedback_enabled, | |
| max_feedback_imgs, | |
| prompt, | |
| neg_prompt, | |
| liked, | |
| disliked, | |
| denoising_steps, | |
| guidance_scale, | |
| feedback_start, | |
| feedback_end, | |
| min_weight, | |
| max_weight, | |
| neg_scale, | |
| batch_size, | |
| seed, | |
| progress=gr.Progress(track_tqdm=True), | |
| ): | |
| try: | |
| if seed < 0: | |
| seed = random.randint(1,9999999999999999) #16 digits is an arbitrary limit | |
| print("seed: ", seed) | |
| max_feedback_imgs = max(0, int(max_feedback_imgs)) | |
| total_images = (len(liked) if liked else 0) + (len(disliked) if disliked else 0) | |
| if not feedback_enabled: | |
| liked = [] | |
| disliked = [] | |
| elif total_images > max_feedback_imgs: | |
| if liked and disliked: | |
| max_disliked = min(len(disliked), max_feedback_imgs // 2) | |
| max_liked = min(len(liked), max_feedback_imgs - max_disliked) | |
| if max_liked > len(liked): | |
| max_disliked = max_feedback_imgs - max_liked | |
| liked = liked[-max_liked:] | |
| disliked = disliked[-max_disliked:] | |
| elif liked: | |
| liked = liked[-max_feedback_imgs:] | |
| disliked = [] | |
| else: | |
| liked = [] | |
| disliked = disliked[-max_feedback_imgs:] | |
| # else: keep all feedback images | |
| generate_kwargs = { | |
| "prompt": prompt, | |
| "negative_prompt": neg_prompt, | |
| "liked": liked, | |
| "disliked": disliked, | |
| "denoising_steps": denoising_steps, | |
| "guidance_scale": guidance_scale, | |
| "feedback_start": feedback_start, | |
| "feedback_end": feedback_end, | |
| "min_weight": min_weight, | |
| "max_weight": max_weight, | |
| "neg_scale": neg_scale, | |
| "seed": seed, | |
| "n_images": batch_size, | |
| } | |
| try: | |
| images = generator.generate(**generate_kwargs) | |
| except RuntimeError as err: | |
| if 'out of memory' in str(err): | |
| generator.reload() | |
| raise | |
| return [(img, f"Image {i+1}") for i, img in enumerate(images)], images, seed | |
| except Exception as err: | |
| raise gr.Error(str(err)) | |
| def add_img_from_list(i, curr_imgs, all_imgs): | |
| if all_imgs is None: | |
| all_imgs = [] | |
| if i >= 0 and i < len(curr_imgs): | |
| all_imgs.append(curr_imgs[i]) | |
| return all_imgs, all_imgs # return (gallery, state) | |
| def add_img(img, all_imgs): | |
| if all_imgs is None: | |
| all_imgs = [] | |
| all_imgs.append(img) | |
| return None, all_imgs, all_imgs | |
| def remove_img_from_list(event: gr.SelectData, imgs): | |
| if event.index >= 0 and event.index < len(imgs): | |
| imgs.pop(event.index) | |
| return imgs, imgs | |
| def duplicate_seed_value(seed): #I don't like the progress bar showing on the previous seed box and this is how I hide it | |
| return seed | |
| with gr.Blocks(css=css) as demo: | |
| liked_imgs = gr.State([]) | |
| disliked_imgs = gr.State([]) | |
| curr_imgs = gr.State([]) | |
| with gr.Row(): | |
| with gr.Column(scale=100): | |
| prompt = gr.Textbox(label="Prompt") | |
| neg_prompt = gr.Textbox(label="Negative prompt", value="lowres, bad anatomy, bad hands, cropped, worst quality") | |
| submit_btn = gr.Button("Generate", variant="primary", min_width="96px") | |
| with gr.Row(equal_height=False): | |
| with gr.Column(): | |
| denoising_steps = gr.Slider(1, 100, value=20, step=1, label="Sampling steps") | |
| guidance_scale = gr.Slider(0.0, 30.0, value=6, step=0.25, label="CFG scale") | |
| batch_size = gr.Slider(1, 10, value=4, step=1, label="Batch size", interactive=False) | |
| seed = gr.Number(-1, minimum=-1, precision=0, label="Seed") | |
| max_feedback_imgs = gr.Slider(0, 20, value=6, step=1, label="Max. feedback images", info="Maximum number of liked/disliked images to be used. If exceeded, only the most recent images will be used as feedback. (NOTE: large number of feedback imgs => high VRAM requirements)") | |
| feedback_enabled = gr.Checkbox(True, label="Enable feedback", interactive=True) | |
| with gr.Accordion("Liked Images", open=True): | |
| liked_img_input = gr.Image(type="pil", shape=(512, 512), height=128, label="Upload liked image") | |
| like_gallery = gr.Gallery(label="👍 Liked images (click to remove)", columns=[3, 4, 3, 4, 5, 6], height=256, allow_preview=False) | |
| clear_liked_btn = gr.Button("Clear likes") | |
| with gr.Accordion("Disliked Images", open=True): | |
| disliked_img_input = gr.Image(type="pil", shape=(512, 512), height=128, label="Upload disliked image") | |
| dislike_gallery = gr.Gallery(label="👎 Disliked images (click to remove)", columns=[3, 4, 3, 4, 5, 6], height=256, allow_preview=False) | |
| clear_disliked_btn = gr.Button("Clear dislikes") | |
| with gr.Accordion("Feedback parameters", open=False): | |
| feedback_start = gr.Slider(0.0, 1.0, value=0.0, label="Feedback start", info="Fraction of denoising steps starting from which to use max. feedback weight.") | |
| feedback_end = gr.Slider(0.0, 1.0, value=0.8, label="Feedback end", info="Up to what fraction of denoising steps to use max. feedback weight.") | |
| feedback_min_weight = gr.Slider(0.0, 1.0, value=0.0, label="Feedback min. weight", info="Attention weight of feedback images when turned off (set to 0.0 to disable)") | |
| feedback_max_weight = gr.Slider(0.0, 1.0, value=0.8, label="Feedback max. weight", info="Attention weight of feedback images when turned on (set to 0.0 to disable)") | |
| feedback_neg_scale = gr.Slider(0.0, 1.0, value=0.5, label="Neg. feedback scale", info="Attention weight of disliked images relative to liked images (set to 0.0 to disable negative feedback)") | |
| with gr.Column(): | |
| gallery = gr.Gallery(label="Generated images") | |
| like_btns = [] | |
| dislike_btns = [] | |
| with gr.Row(): | |
| for i in range(0, 2): | |
| like_btn = gr.Button(f"👍 Image {i+1}", elem_classes="btn-green") | |
| like_btns.append(like_btn) | |
| with gr.Row(): | |
| for i in range(2, 4): | |
| like_btn = gr.Button(f"👍 Image {i+1}", elem_classes="btn-green") | |
| like_btns.append(like_btn) | |
| with gr.Row(): | |
| for i in range(0, 2): | |
| dislike_btn = gr.Button(f"👎 Image {i+1}", elem_classes="btn-red") | |
| dislike_btns.append(dislike_btn) | |
| with gr.Row(): | |
| for i in range(2, 4): | |
| dislike_btn = gr.Button(f"👎 Image {i+1}", elem_classes="btn-red") | |
| dislike_btns.append(dislike_btn) | |
| prev_seed = gr.Number(-1, label="Previous seed", interactive=False) | |
| prev_seed_hid = gr.Number(-1, visible=False) | |
| generate_params = [ | |
| feedback_enabled, | |
| max_feedback_imgs, | |
| prompt, | |
| neg_prompt, | |
| liked_imgs, | |
| disliked_imgs, | |
| denoising_steps, | |
| guidance_scale, | |
| feedback_start, | |
| feedback_end, | |
| feedback_min_weight, | |
| feedback_max_weight, | |
| feedback_neg_scale, | |
| batch_size, | |
| seed, | |
| ] | |
| submit_btn.click(generate_fn, generate_params, [gallery, curr_imgs, prev_seed_hid], queue=True) | |
| prev_seed_hid.change(duplicate_seed_value, prev_seed_hid, prev_seed, queue=False) | |
| for i, like_btn in enumerate(like_btns): | |
| like_btn.click(functools.partial(add_img_from_list, i), [curr_imgs, liked_imgs], [like_gallery, liked_imgs], queue=False) | |
| for i, dislike_btn in enumerate(dislike_btns): | |
| dislike_btn.click(functools.partial(add_img_from_list, i), [curr_imgs, disliked_imgs], [dislike_gallery, disliked_imgs], queue=False) | |
| like_gallery.select(remove_img_from_list, [liked_imgs], [like_gallery, liked_imgs], queue=False) | |
| dislike_gallery.select(remove_img_from_list, [disliked_imgs], [dislike_gallery, disliked_imgs], queue=False) | |
| liked_img_input.upload(add_img, [liked_img_input, liked_imgs], [liked_img_input, like_gallery, liked_imgs], queue=False) | |
| disliked_img_input.upload(add_img, [disliked_img_input, disliked_imgs], [disliked_img_input, dislike_gallery, disliked_imgs], queue=False) | |
| clear_liked_btn.click(lambda: [[], []], None, [liked_imgs, like_gallery], queue=False) | |
| clear_disliked_btn.click(lambda: [[], []], None, [disliked_imgs, dislike_gallery], queue=False) | |
| demo.queue(1) | |
| demo.launch(debug=True) |