Spaces:
Build error
Build error
| import random | |
| import torch | |
| import gc | |
| import gradio as gr | |
| from PIL import Image | |
| from diffusers import DPMSolverMultistepScheduler, StableDiffusionPipeline,\ | |
| StableDiffusionXLPipeline, StableDiffusionUpscalePipeline,\ | |
| DiffusionPipeline | |
| from utils import * | |
| from style import custom_css, beta_header_html | |
| def gen_image(prompt, negative_prompt, width, height, num_steps, | |
| mode, seed, guidance_scale, device): | |
| """ | |
| Run diffusion model to generate image | |
| """ | |
| use_adapter = True | |
| device = f"cuda:{device.split('GPU')[1][1]}" | |
| guidance_scale = float(guidance_scale) | |
| generator = torch.Generator(device).manual_seed(int(seed)) | |
| model_path = DIFFUSION_CHECKPOINTS[mode]["path"] | |
| Text2Image_class = globals()[DIFFUSION_CHECKPOINTS[mode]["pipeline"]] | |
| if DIFFUSION_CHECKPOINTS[mode]["type"] == "pretrained": | |
| pipeline = Text2Image_class.from_pretrained(model_path) | |
| else: | |
| pipeline = Text2Image_class.from_single_file(model_path) | |
| if use_adapter: | |
| if "SDXL 1.0" in mode: | |
| print("Load LoRA model") | |
| pipeline.load_lora_weights("../checkpoints", weight_name="mod2.safetensors") | |
| pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config) | |
| try: | |
| pipeline = pipeline.to(device) | |
| image = pipeline(prompt=prompt, | |
| negative_prompt=negative_prompt, | |
| width=nearest_divisible_by_8(int(width)), | |
| height=nearest_divisible_by_8(int(height)), | |
| num_inference_steps=int(num_steps), | |
| generator=generator, | |
| guidance_scale=guidance_scale).images[0] | |
| del pipeline | |
| torch.cuda.empty_cache() | |
| gc.collect() | |
| except Exception as e: | |
| image = Image.open("stuffs/serverdown.jpg") | |
| print(e) | |
| del pipeline | |
| torch.cuda.empty_cache() | |
| gc.collect() | |
| return image | |
| return image | |
| with gr.Blocks(title="(Beta) TonAI Creative", theme=APP_THEME) as interface1: | |
| gr.HTML(beta_header_html) | |
| with gr.Row(): | |
| with gr.Column(scale=3): | |
| prompt = gr.Textbox(label="Prompt", | |
| placeholder="Tell me what you want to generate", | |
| container=True) | |
| negative_prompt = gr.Textbox(label="Negative Prompt", | |
| placeholder="Instruct the AI model that it should not include", | |
| container=True) | |
| with gr.Row(): | |
| width = gr.Textbox(label="Image Width", value=768) | |
| height = gr.Textbox(label="Image Height", value=768) | |
| with gr.Row(): | |
| seed = gr.Textbox(label="RNG Seed", value=0, scale=1) | |
| guidance_scale = gr.Textbox(label="CFG Scale", value=7, scale=1) | |
| with gr.Row(): | |
| num_steps = gr.components.Slider( | |
| minimum=5, maximum=60, value=20, step=1, | |
| label="Inference Steps" | |
| ) | |
| mode=gr.Dropdown(choices=DIFFUSION_CHECKPOINTS.keys(), label="Mode", | |
| value=list(DIFFUSION_CHECKPOINTS.keys())[1]) | |
| device_choices = display_gpu_info() | |
| device=gr.Dropdown(choices=device_choices, label="Device", value=device_choices[0]) | |
| generate_btn = gr.Button("Generate") | |
| with gr.Column(scale=2): | |
| generate_btn.click( | |
| fn=gen_image, | |
| inputs=[prompt, negative_prompt, width, height, num_steps, mode, seed, guidance_scale, device], | |
| outputs=gr.Image(label="Generated Image", format="png"), | |
| concurrency_limit=5 | |
| ) | |
| interface1.load(lambda: gr.update(value=random.randint(0, 99999)), None, seed) | |
| interface1.load(lambda: gr.update(choices=display_gpu_info(), value=display_gpu_info()[0]), None, device) | |
| # interface = gr.TabbedInterface([interface1, iface2], ["Text-to-text", "image-to-text"]) | |
| allowed_paths=["stuffs/tonai_research_logo.png"] | |
| interface1.queue(default_concurrency_limit=5) | |
| interface1.launch(share=False, | |
| allowed_paths=allowed_paths, | |
| max_threads=5) |