smdbs/CS553_CaseStudy1

1import os2from pathlib import Path3from typing import Tuple, Union4 5import gradio as gr6import numpy as np7 8# import torch9 10# from diffusers import EulerDiscreteScheduler, StableDiffusionPipeline11from huggingface_hub import InferenceClient12from PIL import Image13 14# from prometheus_client import Counter, Summary, start_http_server15 16# # metrics17# REQUEST_COUNTER = Counter("app_requests_total", "Total number of requests")18# LOCAL_COUNTER = Counter("app_local_requests_total", "Total number of local requests")19# API_COUNTER = Counter("app_api_requests_total", "Total number of API requests")20# SUCCESSFUL_REQUESTS = Counter(21#     "app_successful_requests_total", "Total number of successful requests"22# )23# FAILED_REQUESTS = Counter(24#     "app_failed_requests_total", "Total number of failed requests"25# )26# REQUEST_DURATION = Summary(27#     "app_request_duration_seconds", "Time spent processing request"28# )29 30 31MAX_SEED = np.iinfo(np.int32).max32MAX_IMAGE_SIZE = 128033 34# Make sure to set the environment variable HF_TOKEN to your Hugging Face token for using InferenceClient35HF_TOKEN = os.environ.get("HF_TOKEN", "YOUR_HF_TOKEN")36# print(f"Using Hugging Face token: {HF_TOKEN}")37model_name = "stabilityai/stable-diffusion-2-1-base"38client = InferenceClient(model_name, token=HF_TOKEN)39 40scheduler = None41pipe = None42 43 44# def load_model():45#     global scheduler, pipe46#     print("Loading model...")47#     scheduler = EulerDiscreteScheduler.from_pretrained(48#         model_name, subfolder="scheduler"49#     )50#     pipe = StableDiffusionPipeline.from_pretrained(model_name, scheduler=scheduler)51#     pipe.enable_attention_slicing()52#     if torch.cuda.is_available():53#         pipe = pipe.to("cuda")54#     print("Model loaded.")55 56 57def sd_2_1_base(58    prompt: str,59    is_local: bool,60    negative_prompt: str,61    seed: int,62    randomize_seed: bool,63    guidance_scale: float,64    num_inference_steps: int,65    width: int,66    height: int,67    progress: gr.Progress = gr.Progress(track_tqdm=True),68) -> Tuple[Union[np.ndarray, Image.Image, str, Path, None], int]:69    if randomize_seed:70        seed = np.random.randint(0, MAX_SEED)71    # generator = torch.Generator().manual_seed(seed)72 73    # REQUEST_COUNTER.inc()74    if not is_local and (scheduler is None or pipe is None):75        # load_model()76        pass77 78    # with REQUEST_DURATION.time():79    try:80        # if is_local:81        #     LOCAL_COUNTER.inc()82 83        #     image = pipe(84        #         prompt,85        #         negative_prompt=negative_prompt,86        #         guidance_scale=guidance_scale,87        #         num_inference_steps=num_inference_steps,88        #         width=width,89        #         height=height,90        #         generator=generator,91        #     ).images[0]92 93        #     return image, seed94 95        # else:96        # API_COUNTER.inc()97 98        output = client.text_to_image(99            prompt,100            negative_prompt=negative_prompt,101            seed=seed,102            randomize_seed=randomize_seed,103            guidance_scale=guidance_scale,104            num_inference_steps=num_inference_steps,105            width=width,106            height=height,107        )108        return output, seed109    except Exception as e:110        # FAILED_REQUESTS.inc()111        return str(e), seed112 113 114with gr.Blocks() as ui:115    with gr.Column():116        gr.HTML(117            "<h1 style='text-align: center;font-size: 48px;margin-bottom: 24px;'>😎 Cool Image Generator 😎</h1>"118        )119        with gr.Row():120            with gr.Column():121                prompt = gr.Textbox(122                    placeholder="Enter a prompt to generate an image",123                    label="Prompt",124                )125                submit = gr.Button("Generate Image", variant="primary")126                # This part is adopted from Stabilityai's stable-diffusion-3-medium space127                # https://huggingface.co/spaces/stabilityai/stable-diffusion-3-medium128                with gr.Accordion("Advanced Settings", open=False):129                    negative_prompt = gr.Text(130                        label="Negative prompt",131                        max_lines=1,132                        placeholder="Enter a negative prompt",133                    )134 135                    seed = gr.Slider(136                        label="Seed",137                        minimum=0,138                        maximum=MAX_SEED,139                        step=1,140                        value=0,141                    )142 143                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)144 145                    with gr.Row():146 147                        width = gr.Slider(148                            label="Width",149                            minimum=256,150                            maximum=MAX_IMAGE_SIZE,151                            step=64,152                            value=512,153                        )154 155                        height = gr.Slider(156                            label="Height",157                            minimum=256,158                            maximum=MAX_IMAGE_SIZE,159                            step=64,160                            value=512,161                        )162 163                    with gr.Row():164                        guidance_scale = gr.Slider(165                            label="Guidance scale",166                            minimum=0.0,167                            maximum=10.0,168                            step=0.1,169                            value=5.0,170                        )171 172                        num_inference_steps = gr.Slider(173                            label="Number of inference steps",174                            minimum=1,175                            maximum=50,176                            step=1,177                            value=28,178                        )179                gr.Examples(180                    examples=[181                        "Cat in the forest",182                        "Cat flying in the sky",183                        "Cat in the desert",184                    ],185                    inputs=[prompt],186                )187            with gr.Column():188                is_local = gr.Checkbox(189                    label="Check this box to use a local model 🖥️",190                )191                img_out = gr.Image(label="Generated Image")192 193        gr.HTML(194            f"<p style='text-align: center;'>This app uses the <a href='https://huggingface.co/{model_name}'>{model_name}</a> model.</p>"195        )196 197    submit.click(198        sd_2_1_base,199        inputs=[200            prompt,201            is_local,202            negative_prompt,203            seed,204            randomize_seed,205            guidance_scale,206            num_inference_steps,207            width,208            height,209        ],210        outputs=[img_out, seed],211    )212 213if __name__ == "__main__":214    # start_http_server(8000)215    ui.launch()216