OMilosh/DiffusionModelsCourse

1import gradio as gr2import numpy as np3import random4 5# import spaces #[uncomment to use ZeroGPU]6from diffusers import DiffusionPipeline7import torch8 9MAX_SEED = np.iinfo(np.int32).max10MAX_IMAGE_SIZE = 102411 12device = "cuda" if torch.cuda.is_available() else "cpu"13 14available_models = [15    "stabilityai/sdxl-turbo",16    "stabilityai/sd-turbo"]17 18MAX_SEED = np.iinfo(np.int32).max19MAX_IMAGE_SIZE = 102420 21 22def init_model(model_repo_id):23    torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float3224    pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)25    return pipe26 27# @spaces.GPU #[uncomment to use ZeroGPU]28def infer(29    model_repo_id, 30    prompt,31    negative_prompt,32    seed,33    randomize_seed,34    width,35    height,36    guidance_scale,37    num_inference_steps,38    progress=gr.Progress(track_tqdm=True),39):40    pipe = loaded_models[model_repo_id].to(device)41    42    if randomize_seed:43        seed = random.randint(0, MAX_SEED)44 45    generator = torch.Generator().manual_seed(seed)46 47    image = pipe(48        prompt=prompt,49        negative_prompt=negative_prompt,50        guidance_scale=guidance_scale,51        num_inference_steps=num_inference_steps,52        width=width,53        height=height,54        generator=generator,55    ).images[0]56 57    return image, seed58 59 60examples = [61    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",62    "An astronaut riding a green horse",63    "A delicious ceviche cheesecake slice",64]65 66css = """67#col-container {68    margin: 0 auto;69    max-width: 640px;70}71"""72 73with gr.Blocks(css=css) as demo:74    with gr.Column(elem_id="col-container"):75        gr.Markdown(" # Text-to-Image Gradio Template")76 77        with gr.Row():78            prompt = gr.Text(79                label="Prompt",80                show_label=False,81                max_lines=1,82                placeholder="Enter your prompt",83                container=False,84            )85 86            run_button = gr.Button("Run", scale=0, variant="primary")87 88        result = gr.Image(label="Result", show_label=False)89 90        with gr.Accordion("Advanced Settings", open=False):91            model_repo_id = gr.Dropdown(available_models,92                                        value=available_models[0],93                                        multiselect=False,94                                        label="Model",95                                        info="Choose models for generation")96            97            negative_prompt = gr.Text(98                label="Negative prompt",99                max_lines=1,100                placeholder="Enter a negative prompt",101                visible=True,102            )103 104            seed = gr.Slider(105                label="Seed",106                minimum=0,107                maximum=MAX_SEED,108                step=1,109                value=42,110            )111 112            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)113 114            with gr.Row():115                width = gr.Slider(116                    label="Width",117                    minimum=256,118                    maximum=MAX_IMAGE_SIZE,119                    step=32,120                    value=256,  # Replace with defaults that work for your model121                )122 123                height = gr.Slider(124                    label="Height",125                    minimum=256,126                    maximum=MAX_IMAGE_SIZE,127                    step=32,128                    value=256,  # Replace with defaults that work for your model129                )130 131            with gr.Row():132                guidance_scale = gr.Slider(133                    label="Guidance scale",134                    minimum=0.0,135                    maximum=10.0,136                    step=0.1,137                    value=7.0,  # Replace with defaults that work for your model138                )139 140                num_inference_steps = gr.Slider(141                    label="Number of inference steps",142                    minimum=1,143                    maximum=50,144                    step=1,145                    value=20,  # Replace with defaults that work for your model146                )147 148        gr.Examples(examples=examples, inputs=[prompt])149 150        151    loaded_models = {}152 153    for model in available_models:154        loaded_models[model] = init_model(model)155        156    gr.on(157        triggers=[run_button.click, prompt.submit],158        fn=infer,159        inputs=[160            model_repo_id,161            prompt,162            negative_prompt,163            seed,164            randomize_seed,165            width,166            height,167            guidance_scale,168            num_inference_steps,169        ],170        outputs=[result, seed],171    )172 173if __name__ == "__main__":174    demo.launch()175