ClaudeCharest/skprod

1import asyncio2import json3import os4import time5import traceback6 7import gradio as gr8import requests9from gradio_imageslider import ImageSlider10from PIL import Image11 12import tsapi13 14#toggle puts the state we are trying to reach15get_window_url_params = """16    function(url_params) {17        const params = new URLSearchParams(window.location.search);18        url_params = Object.fromEntries(params);19        return url_params;20        }21    """22#load sets of images, upload check the name, depending on name load other images23def main():24    #Retrieve the URL from the environment variable25    url = os.environ.get("TSAPI_URL")26    if url == None:27        url = "https://ypjsritb9e.execute-api.eu-west-2.amazonaws.com"28    #Initialize the client29    def connect(url_params):30        try:31            api_key = url_params['apikey']32            if api_key == None:33                gr.Error("API Key is required")34                return None 35            gr.Info("Connecting to GPU")      36            tsapi.connect(url,api_key)37            gr.Info("GPU is ready")38            return ImageSlider(label="Sketch", type="pil", interactive=True)39        except Exception:40            tsapi.post_bugreport(url, api_key, traceback.format_exc())41            gr.Error("Server maintenance, try again in a few hours")42            return None43 44    def get_tokens(url_params):45        try:46            api_key = url_params['apikey']47            if api_key == None:48                gr.Error("API Key is required")49                return None50            tokens = tsapi.get_tokens(url, api_key)51            return gr.Markdown("Tokens remaining: {}".format(tokens)), gr.JSON({"apikey": api_key}, visible=False, label="URL Params")52        except Exception:53            tsapi.post_bugreport(url, api_key, traceback.format_exc())54            gr.Error("Server maintenance, try again in a few hours")55            return None56    57    def preprocess(url_params, image):58        try:59            api_key = url_params['apikey']60            if api_key == None:61                gr.Error("API Key is required")62                return None63            if image == None:64                return [], ImageSlider(label="Preprocessed Image", type="pil", interactive=True)65            filename = image[0].filename66            images = tsapi.preprocess(url, api_key, filename)67            if images == None:68                gr.Warning("You already have a workload processing, please wait")69                return (None, None)70            return ([image,images[0],images[1]], None)71        except Exception:72            tsapi.post_bugreport(url, api_key, traceback.format_exc())73            gr.Error("Server maintenance, try again in a few hours")74            return None75    def set_preprocessed(preprocessed, control_radio):76        try:77            if preprocessed == None:78                return None, gr.Button("Generate Images", variant="primary", interactive=False),79            if control_radio == "Depth":80                return (Image.open(preprocessed[0][0]),Image.open(preprocessed[1][0])), gr.Button("Generate Images", variant="primary", interactive=True)81            else:82                return (Image.open(preprocessed[0][0]),Image.open(preprocessed[2][0])), gr.Button("Generate Images", variant="primary", interactive=True)83        except Exception:84            gr.Error("Server maintenance, try again in a few hours")85            return None    86    def generate(url_params, sketch, control_type, guidance, pos_prompt, neg_prompt, nbimage):87        try:  88            api_key = os.getenv("TSAPI_KEY")89            if api_key == None and 'apikey' in url_params:90                api_key = url_params['apikey']91                if api_key == None:92                    gr.Error("API Key is required")93                    return None94            gr.Info("Generating Images, current avg  2 minutes")95            sketch = sketch[0].filename96            images = tsapi.sketch_to_image(url, api_key, sketch, control_type, guidance,pos_prompt,neg_prompt,nbimage,None)97            if images == None:98                gr.Warning("You already have a workload processing, please wait")99                return None100        except Exception:101            tsapi.post_bugreport(url, api_key, traceback.format_exc())102            gr.Error("Server maintenance, try again in a few hours")103            return None104        return images105    106    with gr.Blocks(theme=gr.themes.Default(font=[gr.themes.GoogleFont("Plus Jakarta Sans")], primary_hue="purple", secondary_hue="blue")) as demo:107        preprocessed = gr.Gallery(label="Preprocessed Images", columns=2, rows=2, interactive=False,visible=False)108        url_params = gr.JSON({}, visible=False, label="URL Params")109        gr.Markdown(110        """111        """)112        with gr.Row():113            with gr.Column(scale=3):114                with gr.Row():115                    gr.Column(scale=1)116                    with gr.Column(scale=1):117                        image_diff = ImageSlider(label="Preprocessed Image", type="pil", interactive=False)118                    gr.Column(scale=1)119                gallery = gr.Gallery(label="Generated Images", columns=2, rows=2, interactive=False, height="100%")120 121            with gr.Column(scale=1, min_width=300):122                with gr.Row(variant="panel"):123                    with gr.Column():124                        gr.Button("Free Plan", variant="primary", interactive=False)125                        tokens= gr.Markdown("Tokens remaining:  ---")126                tabs = gr.Tabs()127                with tabs:128                    with gr.Tab(label="Sketch to Image", id=0):129                        with gr.Row():130                            with gr.Column(scale=1):131                                pos_prompt = gr.Textbox(label="Positive Prompt",placeholder="Describe what you want to see", interactive=True)132                                nbimage_radio = gr.Radio( choices=[1,2,4],value=4, interactive=True, info="Number of generated images",show_label=False)            133 134                            with gr.Column(scale=1):135                                gr.Markdown(136                                    """137                                    Advanced settings138                                    """)139                                with gr.Row(variant="panel"):140                                    with gr.Column(scale=1):141                                        image_guidance = gr.Image(label="Style reference", type="filepath",  interactive=True)142                                        neg_prompt = gr.Textbox(label="Negative Prompt",placeholder="DO NOT want to see", interactive=True)143                                        control_radio = gr.Radio( label="Preprocessing type", choices=["Depth", "Line"], interactive=True, info="Choose a preprocessed version of your sketch",show_label=True,value = "Depth")            144 145                        generate_button = gr.Button("Generate Images", variant="primary", interactive=False)146        image_diff.upload(fn=preprocess,inputs=[url_params, image_diff], outputs=[preprocessed, image_diff]).then(set_preprocessed,inputs=[preprocessed, control_radio],outputs=[image_diff, generate_button])147        148        control_radio.change(fn=set_preprocessed,inputs=[preprocessed, control_radio],outputs=[image_diff, generate_button])149        150        generate_button.click(fn=generate, inputs=[url_params, image_diff,control_radio, image_guidance, pos_prompt, neg_prompt, nbimage_radio], outputs=gallery).then(get_tokens,inputs=[url_params], outputs=[tokens, url_params])151        demo.load(fn=get_tokens,inputs=[url_params],outputs=[tokens, url_params], js=get_window_url_params).then(fn=connect, inputs=[url_params], outputs=[image_diff])152        demo.launch()153 154 155if __name__ == "__main__":156    asyncio.run(main())157