import torch import numpy as np import gradio as gr import spaces import cv2 import os from typing import Dict from torchvision.transforms.functional import to_tensor, center_crop, resize from PIL import Image from huggingface_hub import Repository token = os.environ['model_fetch'] engine_repo = Repository(local_dir="engine", clone_from="felixrosberg/EngageDiffusion", use_auth_token=token) from engine.ui_model import fetch_model, process_sketch, process_mask, run_model, run_text_to_image engage_logo = Image.open("engage_studios_logo.png").resize((700, 88), Image.Resampling.BICUBIC) engage_logo_mask = np.array(engage_logo.split()[-1])[..., None] / 255 engage_logo_np = np.array(engage_logo.convert('RGB')) pipe = fetch_model() pipe.to('cuda') def inference(user_state, condition_image, settings, prompt, neg_prompt, inference_steps=8, num_images=2, guidance_scale=2.0, guidance_rescale=0.0, enable_freeu=False, height=1024, width=1024, condition_scale=0.5, sketch_detail=1.0, sketch_softness=0.5, inpaint_strength=0.9, exposure=0.0, enable_stylation=False, style_1_down=0.0, style_1_mid=0.0, style_1_up=0.0, style_2_down=0.0, style_2_mid=0.0, style_2_up=0.0, style_3_down=0.0, style_3_mid=0.0, style_3_up=0.0, style_4_down=0.0, style_4_mid=0.0, style_4_up=0.0, seed=None, progress=gr.Progress()): images = run_model(pipe, user_state, condition_image, settings, prompt, neg_prompt, inference_steps=8, num_images=2, guidance_scale, guidance_rescale, enable_freeu, height, width, condition_scale, sketch_detail, sketch_softness, inpaint_strength, exposure, enable_stylation, style_1_down, style_1_mid, style_1_up, style_2_down, style_2_mid, style_2_up, style_3_down, style_3_mid, style_3_up, style_4_down, style_4_mid, style_4_up, seed, progress) for idx, im in enumerate(images): im = np.asarray(im).copy() im[-88:, :700] = im[-88:, :700] * (1 - engage_logo_mask) + engage_logo_np images[idx] = Image.fromarray(np.clip(im.astype('uint8'), 0, 255)) user_state["IMAGE_GALLERY"] += images return user_state["IMAGE_GALLERY"], user_state theme = gr.themes.Base( primary_hue="neutral", radius_size="none", ).set( body_text_color_dark='*neutral_800', embed_radius='*radius_xxs', button_primary_background_fill='*primary_700', button_primary_background_fill_hover='*primary_400', button_primary_background_fill_hover_dark='*primary_400', button_primary_border_color_dark='*primary_200', button_primary_text_color='*primary_50', button_primary_text_color_dark='*primary_50', button_primary_text_color_hover='*primary_50' ) with gr.Blocks(theme=theme) as engage_automotive_lora_demo: session_state = gr.State(value={"IMAGE_GALLERY": [], "SELECTED_IMAGE": None }) diffused_image_out = gr.Gallery(label='Results', show_label=False, columns=[3], rows=[1], object_fit="contain", height="auto", format="png") with gr.Group(): with gr.Row(): prompt_box = gr.Textbox("futuristic dark red car in a white studio", label='Prompt') generate_button = gr.Button("Generate", scale=0) with gr.Row(): settings_dropdown = gr.Dropdown( ["Text to image", "From sketch", "Inpaint", "Inpaint sketch"], value="Text to image", label="Mode", info="Text to image, prompt only. " "From sketch, upload an initial image / sketch in the image editor. " "Inpaint sketch, edits the chosen area of an image. Uses the initial " "image as base for sketches." ) with gr.Accordion("Image Editor", open=False): condition_image = gr.ImageEditor(type='pil', show_label=False, brush=gr.Brush(colors=["#000000"], color_mode="fixed")) with gr.Row(): with gr.Accordion("Settings", open=False): neg_prompt_box = gr.Textbox( "blurry, poor quality, unrealistic", label='Negative Prompt') seed_box = gr.Number(42, label='Seed') inference_steps = gr.Slider(0, 20, value=8, label='Inference Steps', step=1) num_images = gr.Slider(1, 3, value=2, label='Number of Images', step=1) guidance_scale = gr.Slider(0, 10, value=1.5, label='Guidance Scale', step=0.1) guidance_rescale = gr.Slider(0.0, 1.0, value=0.0, label='Guidance Rescale', step=0.1) height = gr.Slider(128, 2048, value=1024, label='Image Height', step=64) width = gr.Slider(128, 2048, value=1024, label='Image Width', step=64) condition_influence = gr.Slider(0.0, 1.0, value=0.5, label='Condition Influence') sketch_detail = gr.Slider(0.0, 1.0, value=0.5, label='Sketch Detail') sketch_softness = gr.Slider(0.0, 1.0, value=0.5, label='Sketch Softness') inpaint_strength = gr.Slider(0.0, 1.0, value=0.8, label='Inpaint Strength') enable_freeu = gr.Checkbox(True, label='FreeU', info='Enables FreeU scaling factors.') with gr.Accordion("Stylation (Experimental)", open=False): with gr.Row(): exposure = gr.Slider(-1.0, 1.0, value=0.0, label='Exposure') enable_stylation = gr.Checkbox(label='Enable Stylation', info='EXPERIMENTAL: We apologize for the ambiguity, ' 'please play around with the sliders to ' 'find a style you like!' 'Warning: Will slow down the generation time.') with gr.Accordion("Style A - Engage Studios Futuristic", open=False): style_A_down = gr.Slider(-1.0, 1.0, value=0.0, label='down') style_A_mid = gr.Slider(-1.0, 1.0, value=0.0, label='mid') style_A_up = gr.Slider(-1.0, 1.0, value=0.0, label='up') with gr.Accordion("Style B - Lighting", open=False): style_B_down = gr.Slider(-1.0, 1.0, value=0.0, label='down') style_B_mid = gr.Slider(-1.0, 1.0, value=0.0, label='mid') style_B_up = gr.Slider(-1.0, 1.0, value=0.0, label='up') with gr.Accordion("Style C - Details A", open=False): style_C_down = gr.Slider(-1.0, 1.0, value=0.0, label='down') style_C_mid = gr.Slider(-1.0, 1.0, value=0.0, label='mid') style_C_up = gr.Slider(-1.0, 1.0, value=0.0, label='up') with gr.Accordion("Style D - Details B", open=False): style_D_down = gr.Slider(-1.0, 1.0, value=0.0, label='down') style_D_mid = gr.Slider(-1.0, 1.0, value=0.0, label='mid') style_D_up = gr.Slider(-1.0, 1.0, value=0.0, label='up') generate_button.click(run_model, inputs=[session_state, condition_image, settings_dropdown, prompt_box, neg_prompt_box, inference_steps, num_images, guidance_scale, guidance_rescale, enable_freeu, height, width, condition_influence, sketch_detail, sketch_softness, inpaint_strength, exposure, enable_stylation, style_A_down, style_A_mid, style_A_up, style_B_down, style_B_mid, style_B_up, style_C_down, style_C_mid, style_C_up, style_D_down, style_D_mid, style_D_up, seed_box], outputs=[diffused_image_out, session_state], show_progress=True) engage_automotive_lora_demo.launch()