ergardt/diffmodels

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