gradio workspace app Jarvis 2024/06/08

app.py

@@ -0,0 +1,423 @@
+import os
+import gc
+import gradio as gr
+import numpy as np
+import torch
+import json
+import spaces
+import config
+import utils
+import logging
+from PIL import Image, PngImagePlugin
+from datetime import datetime
+from diffusers.models import AutoencoderKL
+from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
+
+
+"""
+Changes to base animagine-xl-3.1 log 
+- Cut the wildcard 
+- add in lora pipeline 
+- use let get env variable
+- add in lora strenght variable
+
+"""
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+DESCRIPTION = "Animagine XL 3.1 X Galverse MAMA "
+if not torch.cuda.is_available():
+    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU. </p>"
+IS_COLAB = utils.is_google_colab() or os.getenv("IS_COLAB") == "1"
+HF_TOKEN = os.getenv("HF_TOKEN")
+CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES") == "1"
+MIN_IMAGE_SIZE = 512
+MAX_IMAGE_SIZE = 2048
+USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE") == "1"
+ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1"
+OUTPUT_DIR = os.getenv("OUTPUT_DIR", "./outputs")
+
+MODEL = os.getenv(
+    "MODEL",
+    "https://huggingface.co/cagliostrolab/animagine-xl-3.1/blob/main/animagine-xl-3.1.safetensors",
+)
+LORA_MODEL ="galverse/mama-1.5"
+
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = False
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+
+
+def load_pipeline(model_name):
+    vae = AutoencoderKL.from_pretrained(
+        "madebyollin/sdxl-vae-fp16-fix",
+        torch_dtype=torch.float16,
+    )
+    pipeline = (
+        StableDiffusionXLPipeline.from_single_file
+        if MODEL.endswith(".safetensors")
+        else StableDiffusionXLPipeline.from_pretrained
+    )
+
+    pipe = pipeline(
+        model_name,
+        vae=vae,
+        torch_dtype=torch.float16,
+        custom_pipeline="lpw_stable_diffusion_xl",
+        use_safetensors=True,
+        add_watermarker=False,
+        use_auth_token=HF_TOKEN,
+    )
+    
+    pipe.unload_lora_weights()
+    pipe.load_lora_weights(LORA_MODEL ) #Add in lora to pipe 
+
+    pipe.to(device)
+    return pipe
+
+
+@spaces.GPU
+def generate(
+    prompt: str,
+    negative_prompt: str = "",
+    lora_scale: float = 1.0,  #added in lora strength 
+    seed: int = 0,
+    custom_width: int = 1024,
+    custom_height: int = 1024,
+    guidance_scale: float = 7.0,
+    num_inference_steps: int = 28,
+    sampler: str = "Euler a",
+    aspect_ratio_selector: str = "896 x 1152",
+    style_selector: str = "(None)",
+    quality_selector: str = "Standard v3.1",
+    use_upscaler: bool = False,
+    upscaler_strength: float = 0.55,
+    upscale_by: float = 1.5,
+    add_quality_tags: bool = True,
+    progress=gr.Progress(track_tqdm=True),
+):
+    generator = utils.seed_everything(seed)
+
+    width, height = utils.aspect_ratio_handler(
+        aspect_ratio_selector,
+        custom_width,
+        custom_height,
+    )
+
+  
+
+    prompt, negative_prompt = utils.preprocess_prompt(
+        quality_prompt, quality_selector, prompt, negative_prompt, add_quality_tags
+    )
+    prompt, negative_prompt = utils.preprocess_prompt(
+        styles, style_selector, prompt, negative_prompt
+    )
+
+    width, height = utils.preprocess_image_dimensions(width, height)
+
+    backup_scheduler = pipe.scheduler
+    pipe.scheduler = utils.get_scheduler(pipe.scheduler.config, sampler)
+
+    if use_upscaler:
+        upscaler_pipe = StableDiffusionXLImg2ImgPipeline(**pipe.components)
+    metadata = {
+        "prompt": prompt,
+        "negative_prompt": negative_prompt,
+        "resolution": f"{width} x {height}",
+        "guidance_scale": guidance_scale,
+        "num_inference_steps": num_inference_steps,
+        "seed": seed,
+        "sampler": sampler,
+        "sdxl_style": style_selector,
+        "add_quality_tags": add_quality_tags,
+        "quality_tags": quality_selector,
+    }
+    #add in lora 
+    metadata["cross_attention_kwargs"] = {"scale": lora_scale}
+
+    if use_upscaler:
+        new_width = int(width * upscale_by)
+        new_height = int(height * upscale_by)
+        metadata["use_upscaler"] = {
+            "upscale_method": "nearest-exact",
+            "upscaler_strength": upscaler_strength,
+            "upscale_by": upscale_by,
+            "new_resolution": f"{new_width} x {new_height}",
+        }
+    else:
+        metadata["use_upscaler"] = None
+        metadata["Model"] = {
+            "Model": DESCRIPTION,
+            "Model hash": "e3c47aedb0",
+        }
+    
+    logger.info(json.dumps(metadata, indent=4))
+
+    try:
+        if use_upscaler:
+            latents = pipe(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                width=width,
+                height=height,
+                guidance_scale=guidance_scale,
+                num_inference_steps=num_inference_steps,
+                generator=generator,
+                cross_attention_kwargs={"scale":  lora_scale}, #add in lora scale setting 
+                output_type="latent",
+            ).images
+            upscaled_latents = utils.upscale(latents, "nearest-exact", upscale_by)
+            images = upscaler_pipe(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                image=upscaled_latents,
+                guidance_scale=guidance_scale,
+                num_inference_steps=num_inference_steps,
+                strength=upscaler_strength,
+                generator=generator,
+                output_type="pil",
+            ).images
+        else:
+            images = pipe(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                width=width,
+                height=height,
+                guidance_scale=guidance_scale,
+                num_inference_steps=num_inference_steps,
+                generator=generator,
+                cross_attention_kwargs={"scale":  lora_scale}, #add in lora scale setting 
+                output_type="pil",
+            ).images
+
+        if images:
+            image_paths = [
+                utils.save_image(image, metadata, OUTPUT_DIR, IS_COLAB)
+                for image in images
+            ]
+
+            for image_path in image_paths:
+                logger.info(f"Image saved as {image_path} with metadata")
+
+        return image_paths, metadata
+    except Exception as e:
+        logger.exception(f"An error occurred: {e}")
+        raise
+    finally:
+        if use_upscaler:
+            del upscaler_pipe
+        pipe.scheduler = backup_scheduler
+        utils.free_memory()
+
+
+if torch.cuda.is_available():
+    pipe = load_pipeline(MODEL)
+    logger.info("Loaded on Device!")
+else:
+    pipe = None
+
+styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in config.style_list}
+quality_prompt = {
+    k["name"]: (k["prompt"], k["negative_prompt"]) for k in config.quality_prompt_list
+}
+
+
+
+with gr.Blocks(css="style.css", theme="NoCrypt/miku@1.2.1") as demo:
+    title = gr.HTML(
+        f"""<h1><span>{DESCRIPTION}</span></h1>""",
+        elem_id="title",
+    )
+    gr.Markdown(
+        f"""Gradio demo for [cagliostrolab/animagine-xl-3.1](https://huggingface.co/cagliostrolab/animagine-xl-3.1)""",
+        elem_id="subtitle",
+    )
+    gr.DuplicateButton(
+        value="Duplicate Space for private use",
+        elem_id="duplicate-button",
+        visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
+    )
+    with gr.Row():
+        with gr.Column(scale=2):
+            with gr.Tab("Txt2img"):
+                with gr.Group():
+                    prompt = gr.Text(
+                        label="Prompt",
+                        max_lines=5,
+                        placeholder="Enter your prompt",
+                    )
+                    negative_prompt = gr.Text(
+                        label="Negative Prompt",
+                        max_lines=5,
+                        placeholder="Enter a negative prompt",
+                    )
+                    lora_scale = gr.Slider(
+                            label="lora strength",
+                            minimum=0,
+                            maximum=1.2,
+                            step=0.1,
+                            value=1,
+                    )
+                    with gr.Accordion(label="Quality Tags", open=True):
+                        add_quality_tags = gr.Checkbox(
+                            label="Add Quality Tags", value=True
+                        )
+                        quality_selector = gr.Dropdown(
+                            label="Quality Tags Presets",
+                            interactive=True,
+                            choices=list(quality_prompt.keys()),
+                            value="Standard v3.1",
+                        )
+            with gr.Tab("Advanced Settings"):
+                with gr.Group():
+                    style_selector = gr.Radio(
+                        label="Style Preset",
+                        container=True,
+                        interactive=True,
+                        choices=list(styles.keys()),
+                        value="(None)",
+                    )
+                with gr.Group():
+                    aspect_ratio_selector = gr.Radio(
+                        label="Aspect Ratio",
+                        choices=config.aspect_ratios,
+                        value="896 x 1152",
+                        container=True,
+                    )
+                with gr.Group(visible=False) as custom_resolution:
+                    with gr.Row():
+                        custom_width = gr.Slider(
+                            label="Width",
+                            minimum=MIN_IMAGE_SIZE,
+                            maximum=MAX_IMAGE_SIZE,
+                            step=8,
+                            value=1024,
+                        )
+                        custom_height = gr.Slider(
+                            label="Height",
+                            minimum=MIN_IMAGE_SIZE,
+                            maximum=MAX_IMAGE_SIZE,
+                            step=8,
+                            value=1024,
+                        )
+                with gr.Group():
+                    use_upscaler = gr.Checkbox(label="Use Upscaler", value=False)
+                    with gr.Row() as upscaler_row:
+                        upscaler_strength = gr.Slider(
+                            label="Strength",
+                            minimum=0,
+                            maximum=1,
+                            step=0.05,
+                            value=0.55,
+                            visible=False,
+                        )
+                        upscale_by = gr.Slider(
+                            label="Upscale by",
+                            minimum=1,
+                            maximum=1.5,
+                            step=0.1,
+                            value=1.5,
+                            visible=False,
+                        )
+                with gr.Group():
+                    sampler = gr.Dropdown(
+                        label="Sampler",
+                        choices=config.sampler_list,
+                        interactive=True,
+                        value="Euler a",
+                    )
+                with gr.Group():
+                    seed = gr.Slider(
+                        label="Seed", minimum=0, maximum=utils.MAX_SEED, step=1, value=0
+                    )
+                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+                with gr.Group():
+                    with gr.Row():
+                        guidance_scale = gr.Slider(
+                            label="Guidance scale",
+                            minimum=1,
+                            maximum=12,
+                            step=0.1,
+                            value=7.0,
+                        )
+                        num_inference_steps = gr.Slider(
+                            label="Number of inference steps",
+                            minimum=1,
+                            maximum=50,
+                            step=1,
+                            value=28,
+                        )
+        with gr.Column(scale=3):
+            with gr.Blocks():
+                run_button = gr.Button("Generate", variant="primary")
+            result = gr.Gallery(
+                label="Result", 
+                columns=1, 
+                height='100%', 
+                preview=True, 
+                show_label=False
+            )
+            with gr.Accordion(label="Generation Parameters", open=False):
+                gr_metadata = gr.JSON(label="metadata", show_label=False)
+            gr.Examples(
+                examples=config.examples,
+                inputs=prompt,
+                outputs=[result, gr_metadata],
+                fn=lambda *args, **kwargs: generate(*args, use_upscaler=True, **kwargs),
+                cache_examples=CACHE_EXAMPLES,
+            )
+    use_upscaler.change(
+        fn=lambda x: [gr.update(visible=x), gr.update(visible=x)],
+        inputs=use_upscaler,
+        outputs=[upscaler_strength, upscale_by],
+        queue=False,
+        api_name=False,
+    )
+    aspect_ratio_selector.change(
+        fn=lambda x: gr.update(visible=x == "Custom"),
+        inputs=aspect_ratio_selector,
+        outputs=custom_resolution,
+        queue=False,
+        api_name=False,
+    )
+
+    gr.on(
+        triggers=[
+            prompt.submit,
+            negative_prompt.submit,
+            run_button.click,
+        ],
+        fn=utils.randomize_seed_fn,
+        inputs=[seed, randomize_seed],
+        outputs=seed,
+        queue=False,
+        api_name=False,
+    ).then(
+        fn=generate,
+        inputs=[
+            prompt,
+            negative_prompt,
+            lora_scale,
+            seed,
+            custom_width,
+            custom_height,
+            guidance_scale,
+            num_inference_steps,
+            sampler,
+            aspect_ratio_selector,
+            style_selector,
+            quality_selector,
+            use_upscaler,
+            upscaler_strength,
+            upscale_by,
+            add_quality_tags,
+        ],
+        outputs=[result, gr_metadata],
+        api_name="run",
+    )
+
+if __name__ == "__main__":
+    demo.queue(max_size=20).launch(debug=IS_COLAB, share=IS_COLAB)

config.py

@@ -0,0 +1,111 @@
+examples = [
+    "1girl, souryuu asuka langley, neon genesis evangelion, plugsuit, pilot suit, red bodysuit, sitting, crossing legs, black eye patch, cat hat, throne, symmetrical, looking down, from bottom, looking at viewer, outdoors",
+    "1boy, male focus, yuuki makoto \(persona 3\), persona 3, black jacket, white shirt, long sleeves, closed mouth, glowing eyes, gun, hair over one eye, holding gun, handgun, looking at viewer, solo, upper body",
+    "1girl, makima \(chainsaw man\), chainsaw man, black jacket, black necktie, black pants, braid, business suit, fingernails, formal, hand on own chin, jacket on shoulders, light smile, long sleeves, looking at viewer, looking up, medium breasts, office lady, smile, solo, suit, upper body, white shirt, outdoors",
+    "1boy, male focus, gojou satoru, jujutsu kaisen, black jacket, blindfold lift, blue eyes, glowing, glowing eyes, high collar, jacket, jujutsu tech uniform, solo, grin, white hair",
+    "1girl, cagliostro, granblue fantasy, violet eyes, standing, hand on own chin, looking at object, smile, closed mouth, table, beaker, glass tube, experiment apparatus, dark room, laboratory",
+    "kimi no na wa., building, cityscape, cloud, cloudy sky, gradient sky, lens flare, no humans, outdoors, power lines, scenery, shooting star, sky, sparkle, star \(sky\), starry sky, sunset, tree, utility pole",
+]
+
+quality_prompt_list = [
+    {
+        "name": "(None)",
+        "prompt": "{prompt}",
+        "negative_prompt": "nsfw, lowres",
+    },
+    {
+        "name": "Standard v3.0",
+        "prompt": "{prompt}, masterpiece, best quality",
+        "negative_prompt": "nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, artist name",
+    },
+    {
+        "name": "Standard v3.1",
+        "prompt": "{prompt}, masterpiece, best quality, very aesthetic, absurdres",
+        "negative_prompt": "nsfw, lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts, low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract]",
+    },
+    {
+        "name": "Light v3.1",
+        "prompt": "{prompt}, (masterpiece), best quality, very aesthetic, perfect face",
+        "negative_prompt": "nsfw, (low quality, worst quality:1.2), very displeasing, 3d, watermark, signature, ugly, poorly drawn",
+    },
+    {
+        "name": "Heavy v3.1",
+        "prompt": "{prompt}, (masterpiece), (best quality), (ultra-detailed), very aesthetic, illustration, disheveled hair, perfect composition, moist skin, intricate details",
+        "negative_prompt": "nsfw, longbody, lowres, bad anatomy, bad hands, missing fingers, pubic hair, extra digit, fewer digits, cropped, worst quality, low quality, very displeasing",
+    },
+]
+
+sampler_list = [
+    "DPM++ 2M Karras",
+    "DPM++ SDE Karras",
+    "DPM++ 2M SDE Karras",
+    "Euler",
+    "Euler a",
+    "DDIM",
+]
+
+aspect_ratios = [
+    "1024 x 1024",
+    "1152 x 896",
+    "896 x 1152",
+    "1216 x 832",
+    "832 x 1216",
+    "1344 x 768",
+    "768 x 1344",
+    "1536 x 640",
+    "640 x 1536",
+    "Custom",
+]
+
+style_list = [
+    {
+        "name": "(None)",
+        "prompt": "{prompt}",
+        "negative_prompt": "",
+    },
+    {
+        "name": "Cinematic",
+        "prompt": "{prompt}, cinematic still, emotional, harmonious, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy",
+        "negative_prompt": "nsfw, cartoon, graphic, text, painting, crayon, graphite, abstract, glitch, deformed, mutated, ugly, disfigured",
+    },
+    {
+        "name": "Photographic",
+        "prompt": "{prompt}, cinematic photo, 35mm photograph, film, bokeh, professional, 4k, highly detailed",
+        "negative_prompt": "nsfw, drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly",
+    },
+    {
+        "name": "Anime",
+        "prompt": "{prompt}, anime artwork, anime style, key visual, vibrant, studio anime, highly detailed",
+        "negative_prompt": "nsfw, photo, deformed, black and white, realism, disfigured, low contrast",
+    },
+    {
+        "name": "Manga",
+        "prompt": "{prompt}, manga style, vibrant, high-energy, detailed, iconic, Japanese comic style",
+        "negative_prompt": "nsfw, ugly, deformed, noisy, blurry, low contrast, realism, photorealistic, Western comic style",
+    },
+    {
+        "name": "Digital Art",
+        "prompt": "{prompt}, concept art, digital artwork, illustrative, painterly, matte painting, highly detailed",
+        "negative_prompt": "nsfw, photo, photorealistic, realism, ugly",
+    },
+    {
+        "name": "Pixel art",
+        "prompt": "{prompt}, pixel-art, low-res, blocky, pixel art style, 8-bit graphics",
+        "negative_prompt": "nsfw, sloppy, messy, blurry, noisy, highly detailed, ultra textured, photo, realistic",
+    },
+    {
+        "name": "Fantasy art",
+        "prompt": "{prompt}, ethereal fantasy concept art, magnificent, celestial, ethereal, painterly, epic, majestic, magical, fantasy art, cover art, dreamy",
+        "negative_prompt": "nsfw, photographic, realistic, realism, 35mm film, dslr, cropped, frame, text, deformed, glitch, noise, noisy, off-center, deformed, cross-eyed, closed eyes, bad anatomy, ugly, disfigured, sloppy, duplicate, mutated, black and white",
+    },
+    {
+        "name": "Neonpunk",
+        "prompt": "{prompt}, neonpunk style, cyberpunk, vaporwave, neon, vibes, vibrant, stunningly beautiful, crisp, detailed, sleek, ultramodern, magenta highlights, dark purple shadows, high contrast, cinematic, ultra detailed, intricate, professional",
+        "negative_prompt": "nsfw, painting, drawing, illustration, glitch, deformed, mutated, cross-eyed, ugly, disfigured",
+    },
+    {
+        "name": "3D Model",
+        "prompt": "{prompt}, professional 3d model, octane render, highly detailed, volumetric, dramatic lighting",
+        "negative_prompt": "nsfw, ugly, deformed, noisy, low poly, blurry, painting",
+    },
+]

requirements.txt

@@ -0,0 +1,10 @@
+accelerate==0.27.2
+diffusers==0.26.3
+gradio==4.20.0
+invisible-watermark==0.2.0
+Pillow==10.2.0
+spaces==0.24.0
+torch==2.0.1
+transformers==4.38.1
+omegaconf==2.3.0
+timm==0.9.10

style.css

@@ -0,0 +1,63 @@
+:root {
+    --title-font-size: clamp(1.5rem, 6vw, 3rem);
+    --subtitle-font-size: clamp(1rem, 2vw, 1.2rem);
+  }
+  
+  h1 {
+    text-align: center;
+    font-size: var(--title-font-size);
+    display: block;
+  }
+  
+  h2 {
+    text-align: center;
+    font-size: 2rem;
+    display: block;
+  }
+  
+  #duplicate-button {
+    display: block;
+    margin: 1rem auto;
+    color: #fff;
+    background: #1565c0;
+    border-radius: 100vh;
+    padding: 0.5rem 1rem;
+  }
+  
+  #component-0 {
+    max-width: 85%;
+    margin: 2rem auto;
+    padding: 2rem;
+  }
+  
+  @media (max-width: 600px) {
+    #component-0 {
+      max-width: 100%;
+      padding: 0.5rem;
+    }
+  }
+  
+  #title-container {
+    text-align: center;
+    padding: 2rem 0;
+  }
+  
+  #title {
+    font-size: var(--title-font-size);
+    color: #333;
+    font-family: 'Helvetica Neue', sans-serif;
+    text-transform: uppercase;
+    background: transparent;
+  }
+  
+  #title span {
+    background: linear-gradient(45deg, #4EACEF, #28b485);
+    background-clip: text;
+    color: transparent;
+  }
+  
+  #subtitle {
+    text-align: center;
+    font-size: var(--subtitle-font-size);
+    margin-top: 1rem;
+  }

utils.py

@@ -0,0 +1,164 @@
+import gc
+import os
+import random
+import numpy as np
+import json
+import torch
+import uuid
+from PIL import Image, PngImagePlugin
+from datetime import datetime
+from dataclasses import dataclass
+from typing import Callable, Dict, Optional, Tuple
+from diffusers import (
+    DDIMScheduler,
+    DPMSolverMultistepScheduler,
+    DPMSolverSinglestepScheduler,
+    EulerAncestralDiscreteScheduler,
+    EulerDiscreteScheduler,
+)
+
+MAX_SEED = np.iinfo(np.int32).max
+
+
+@dataclass
+class StyleConfig:
+    prompt: str
+    negative_prompt: str
+
+
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
+
+
+def seed_everything(seed: int) -> torch.Generator:
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    generator = torch.Generator()
+    generator.manual_seed(seed)
+    return generator
+
+
+def parse_aspect_ratio(aspect_ratio: str) -> Optional[Tuple[int, int]]:
+    if aspect_ratio == "Custom":
+        return None
+    width, height = aspect_ratio.split(" x ")
+    return int(width), int(height)
+
+
+def aspect_ratio_handler(
+    aspect_ratio: str, custom_width: int, custom_height: int
+) -> Tuple[int, int]:
+    if aspect_ratio == "Custom":
+        return custom_width, custom_height
+    else:
+        width, height = parse_aspect_ratio(aspect_ratio)
+        return width, height
+
+
+def get_scheduler(scheduler_config: Dict, name: str) -> Optional[Callable]:
+    scheduler_factory_map = {
+        "DPM++ 2M Karras": lambda: DPMSolverMultistepScheduler.from_config(
+            scheduler_config, use_karras_sigmas=True
+        ),
+        "DPM++ SDE Karras": lambda: DPMSolverSinglestepScheduler.from_config(
+            scheduler_config, use_karras_sigmas=True
+        ),
+        "DPM++ 2M SDE Karras": lambda: DPMSolverMultistepScheduler.from_config(
+            scheduler_config, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++"
+        ),
+        "Euler": lambda: EulerDiscreteScheduler.from_config(scheduler_config),
+        "Euler a": lambda: EulerAncestralDiscreteScheduler.from_config(
+            scheduler_config
+        ),
+        "DDIM": lambda: DDIMScheduler.from_config(scheduler_config),
+    }
+    return scheduler_factory_map.get(name, lambda: None)()
+
+
+def free_memory() -> None:
+    torch.cuda.empty_cache()
+    gc.collect()
+
+
+def preprocess_prompt(
+    style_dict,
+    style_name: str,
+    positive: str,
+    negative: str = "",
+    add_style: bool = True,
+) -> Tuple[str, str]:
+    p, n = style_dict.get(style_name, style_dict["(None)"])
+
+    if add_style and positive.strip():
+        formatted_positive = p.format(prompt=positive)
+    else:
+        formatted_positive = positive
+
+    combined_negative = n
+    if negative.strip():
+        if combined_negative:
+            combined_negative += ", " + negative
+        else:
+            combined_negative = negative
+
+    return formatted_positive, combined_negative
+
+
+def common_upscale(
+    samples: torch.Tensor,
+    width: int,
+    height: int,
+    upscale_method: str,
+) -> torch.Tensor:
+    return torch.nn.functional.interpolate(
+        samples, size=(height, width), mode=upscale_method
+    )
+
+
+def upscale(
+    samples: torch.Tensor, upscale_method: str, scale_by: float
+) -> torch.Tensor:
+    width = round(samples.shape[3] * scale_by)
+    height = round(samples.shape[2] * scale_by)
+    return common_upscale(samples, width, height, upscale_method)
+
+
+
+def get_random_line_from_file(file_path: str) -> str:
+    with open(file_path, "r") as file:
+        lines = file.readlines()
+        if not lines:
+            return ""
+        return random.choice(lines).strip()
+
+def preprocess_image_dimensions(width, height):
+    if width % 8 != 0:
+        width = width - (width % 8)
+    if height % 8 != 0:
+        height = height - (height % 8)
+    return width, height
+
+
+def save_image(image, metadata, output_dir, is_colab):
+    if is_colab:
+        current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
+        filename = f"image_{current_time}.png"   
+    else:
+        filename = str(uuid.uuid4()) + ".png"
+    os.makedirs(output_dir, exist_ok=True)
+    filepath = os.path.join(output_dir, filename)
+    metadata_str = json.dumps(metadata)
+    info = PngImagePlugin.PngInfo()
+    info.add_text("metadata", metadata_str)
+    image.save(filepath, "PNG", pnginfo=info)
+    return filepath
+    
+    
+def is_google_colab():
+    try:
+        import google.colab
+        return True
+    except: