Update app.py

app.py

@@ -1,152 +1,97 @@
+from share import *
+import config
+
+import cv2
+import einops
 import gradio as gr
-from random import randint, sample
-from all_models import models
-import csv
-import os
-
-# Assuming you have a function to calculate ELO ratings
-def init_model_scores(file_path='model_scores.csv'):
-    # Check if the CSV file exists, if not, create it with headers
-    if not os.path.isfile(csv_file_path):
-        with open(csv_file_path, 'w', newline='') as file:
-            writer = csv.writer(file)
-            writer.writerow(["Model Name", "Score"])
-            for model in models:
-                # make a entry for each model
-                writer.writerow([model, 0])
-    
-def update_elo_ratings(user_vote, csv_file_path='model_scores.csv'):
-    # Logic to update ELO ratings based on user vote
-    
-    # Read the current scores from the CSV file
-    scores = {}
-    with open(csv_file_path, 'r') as file:
-        reader = csv.reader(file)
-        next(reader) # Skip the header row
-        for row in reader:
-            scores[row[0]] = int(row[1])
-    
-    # Update the score for the selected model
-    if user_vote in scores:
-        scores[user_vote] += 1 # Increment the score
-    else:
-        scores[user_vote] = 1 # Add the model with a score of 1
-    
-    # Write the updated scores back to the CSV file
-    with open(csv_file_path, 'w', newline='') as file:
-        writer = csv.writer(file)
-        writer.writerow(["Model Name", "Score"]) # Write the header row
-        for model, score in scores.items():
-            writer.writerow([model, score])
-
-
-# Function to compare two models
-def compare_models(prompt):
-    model1, model2 = sample(models, 2)
-    image1, model_name1 = gen_fn(model1, prompt)
-    image2, model_name2 = gen_fn(model2, prompt)
-    return image1, model_name1, image2, model_name2
-
-# User voting logic
-def handle_vote(user_vote):
-    init_model_scores()
-    # Assuming user_vote is a string indicating the preferred model
-    # Update ELO ratings based on user vote
-    update_elo_ratings(user_vote)
-    
-# Leaderboard display logic
-def display_leaderboard():
-    # Logic to display leaderboard based on ELO ratings
-    pass
-
-# Your existing Gradio setup code here...
-def load_fn(models):
-    global models_load
-    models_load = {}
-    
-    for model in models:
-        if model not in models_load.keys():
-            try:
-                m = gr.load(f'models/{model}')
-            except Exception as error:
-                m = gr.Interface(lambda txt: None, ['text'], ['image'])
-            models_load.update({model: m})
-
-
-load_fn(models)
-
-
-num_models = 6
-default_models = models[:num_models]
-
-
-def extend_choices(choices):
-    return choices + (num_models - len(choices)) * ['NA']
-
-
-def update_imgbox(choices):
-    choices_plus = extend_choices(choices)
-    return [gr.Image(None, label = m, visible = (m != 'NA')) for m in choices_plus]
-
-    
-def gen_fn(model_str, prompt):
-    if model_str == 'NA':
-        return None
-    noise = str(randint(0, 99999999999))
-    return models_load[model_str](f'{prompt} {noise}')
-    
-# Modified gen_fn function to return model name
-def gen_fn(model_str, prompt):
-    if model_str == 'NA':
-        return None, None
-    noise = str(randint(0, 99999999999))
-    image = models_load[model_str](f'{prompt} {noise}')
-    return image, model_str
-
-
-with gr.Blocks() as ImageGenarationArena:
-    with gr.Column('model A', variant='panel', width=2, height=150) as col:
-    #with gr.Tab('model B'):
-        model_choice2 = gr.Dropdown(models, label = 'Choose model', value = models[0], filterable = False)
-        txt_input2 = gr.Textbox(label = 'Prompt text')
-        
-        max_images = 6
-        num_images = gr.Slider(1, max_images, value = max_images, step = 1, label = 'Number of images')
-        
-        gen_button2 = gr.Button('Generate')
-        stop_button2 = gr.Button('Stop', variant = 'secondary', interactive = False)
-        gen_button2.click(lambda s: gr.update(interactive = True), None, stop_button2)
-        
-        with gr.Row():
-            output2 = [gr.Image(label = '') for _ in range(max_images)]
-
-        for i, o in enumerate(output2):
-            img_i = gr.Number(i, visible = False)
-            num_images.change(lambda i, n: gr.update(visible = (i < n)), [img_i, num_images], o)
-            gen_event2 = gen_button2.click(lambda i, n, m, t: gen_fn(m, t) if (i < n) else None, [img_i, num_images, model_choice2, txt_input2], o)
-            stop_button2.click(lambda s: gr.update(interactive = False), None, stop_button2, cancels = [gen_event2])
-
-    with gr.Column('model B', variant='panel', width=2, height=150) as col:    
-    #with gr.Tab('model A'):
-        model_choice2 = gr.Dropdown(models, label = 'Choose model', value = models[0], filterable = False)
-        txt_input2 = gr.Textbox(label = 'Prompt text')
-        
-        max_images = 6
-        num_images = gr.Slider(1, max_images, value = max_images, step = 1, label = 'Number of images')
-        
-        gen_button2 = gr.Button('Generate')
-        stop_button2 = gr.Button('Stop', variant = 'secondary', interactive = False)
-        gen_button2.click(lambda s: gr.update(interactive = True), None, stop_button2)
-        
-        with gr.Row():
-            output2 = [gr.Image(label = '') for _ in range(max_images)]
-
-        for i, o in enumerate(output2):
-            img_i = gr.Number(i, visible = False)
-            num_images.change(lambda i, n: gr.update(visible = (i < n)), [img_i, num_images], o)
-            gen_event2 = gen_button2.click(lambda i, n, m, t: gen_fn(m, t) if (i < n) else None, [img_i, num_images, model_choice2, txt_input2], o)
-            stop_button2.click(lambda s: gr.update(interactive = False), None, stop_button2, cancels = [gen_event2])
-
-    
-ImageGenarationArena.queue(concurrency_count = 36)                        
-ImageGenarationArena.launch()
+import numpy as np
+import torch
+import random
+
+from pytorch_lightning import seed_everything
+from annotator.util import resize_image, HWC3
+from annotator.uniformer import UniformerDetector
+from cldm.model import create_model, load_state_dict
+from cldm.ddim_hacked import DDIMSampler
+
+
+apply_uniformer = UniformerDetector()
+
+model = create_model('./models/cldm_v15.yaml').cpu()
+model.load_state_dict(load_state_dict('./models/control_sd15_seg.pth', location='cuda'))
+model = model.cuda()
+ddim_sampler = DDIMSampler(model)
+
+
+def process(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, detect_resolution, ddim_steps, guess_mode, strength, scale, seed, eta):
+    with torch.no_grad():
+        input_image = HWC3(input_image)
+        detected_map = apply_uniformer(resize_image(input_image, detect_resolution))
+        img = resize_image(input_image, image_resolution)
+        H, W, C = img.shape
+
+        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_NEAREST)
+
+        control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
+        control = torch.stack([control for _ in range(num_samples)], dim=0)
+        control = einops.rearrange(control, 'b h w c -> b c h w').clone()
+
+        if seed == -1:
+            seed = random.randint(0, 65535)
+        seed_everything(seed)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=False)
+
+        cond = {"c_concat": [control], "c_crossattn": [model.get_learned_conditioning([prompt + ', ' + a_prompt] * num_samples)]}
+        un_cond = {"c_concat": None if guess_mode else [control], "c_crossattn": [model.get_learned_conditioning([n_prompt] * num_samples)]}
+        shape = (4, H // 8, W // 8)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=True)
+
+        model.control_scales = [strength * (0.825 ** float(12 - i)) for i in range(13)] if guess_mode else ([strength] * 13)  # Magic number. IDK why. Perhaps because 0.825**12<0.01 but 0.826**12>0.01
+        samples, intermediates = ddim_sampler.sample(ddim_steps, num_samples,
+                                                     shape, cond, verbose=False, eta=eta,
+                                                     unconditional_guidance_scale=scale,
+                                                     unconditional_conditioning=un_cond)
+
+        if config.save_memory:
+            model.low_vram_shift(is_diffusing=False)
+
+        x_samples = model.decode_first_stage(samples)
+        x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
+
+        results = [x_samples[i] for i in range(num_samples)]
+    return [detected_map] + results
+
+
+block = gr.Blocks().queue()
+with block:
+    with gr.Row():
+        gr.Markdown("## Control Stable Diffusion with Segmentation Maps")
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(source='upload', type="numpy")
+            prompt = gr.Textbox(label="Prompt")
+            run_button = gr.Button(label="Run")
+            with gr.Accordion("Advanced options", open=False):
+                num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1)
+                image_resolution = gr.Slider(label="Image Resolution", minimum=256, maximum=768, value=512, step=64)
+                strength = gr.Slider(label="Control Strength", minimum=0.0, maximum=2.0, value=1.0, step=0.01)
+                guess_mode = gr.Checkbox(label='Guess Mode', value=False)
+                detect_resolution = gr.Slider(label="Segmentation Resolution", minimum=128, maximum=1024, value=512, step=1)
+                ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1)
+                scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
+                seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, randomize=True)
+                eta = gr.Number(label="eta (DDIM)", value=0.0)
+                a_prompt = gr.Textbox(label="Added Prompt", value='best quality, extremely detailed')
+                n_prompt = gr.Textbox(label="Negative Prompt",
+                                      value='longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality')
+        with gr.Column():
+            result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2, height='auto')
+    ips = [input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, detect_resolution, ddim_steps, guess_mode, strength, scale, seed, eta]
+    run_button.click(fn=process, inputs=ips, outputs=[result_gallery])
+
+
+block.launch(server_name='0.0.0.0')