Upload 4 files

README.md

@@ -4,7 +4,7 @@ emoji: 🪈
 colorFrom: purple
 colorTo: yellow
 sdk: gradio
-sdk_version: 5.4.0
+sdk_version: 5.14.0
 app_file: app.py
 pinned: false
 license: other

app.py

@@ -1,126 +1,89 @@
-import os
-import spaces
-import torch
-from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
-import gradio as gr
-import random
-import tqdm
-
-# Enable TQDM progress tracking
-tqdm.monitor_interval = 0
-
-# Load the diffusion pipeline
-pipe = StableDiffusionXLPipeline.from_pretrained(
-    "kayfahaarukku/UrangDiffusion-2.0", 
-    torch_dtype=torch.float16, 
-    custom_pipeline="lpw_stable_diffusion_xl",
-)
-pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
-
-# Function to generate an image
-@spaces.GPU  # Adjust the duration as needed
-def generate_image(prompt, negative_prompt, use_defaults, resolution, guidance_scale, num_inference_steps, seed, randomize_seed, progress=gr.Progress()):
-    pipe.to('cuda')  # Move the model to GPU when the function is called
-    
-    if randomize_seed:
-        seed = random.randint(0, 99999999)
-    if use_defaults:
-        prompt = f"{prompt}, best quality, amazing quality, very aesthetic"
-        negative_prompt = f"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], {negative_prompt}"
-    generator = torch.manual_seed(seed)
-    
-    def callback(step, timestep, latents):
-        progress(step / num_inference_steps)
-        return
-    
-    width, height = map(int, resolution.split('x'))
-    image = pipe(
-        prompt, 
-        negative_prompt=negative_prompt,
-        width=width,
-        height=height, 
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        generator=generator,
-        callback=callback,
-        callback_steps=1
-    ).images[0]
-
-    torch.cuda.empty_cache()
-
-    metadata_text = f"{prompt}\nNegative prompt: {negative_prompt}\nSteps: {num_inference_steps}, Sampler: Euler a, Size: {width}x{height}, Seed: {seed}, CFG scale: {guidance_scale}"
-
-    return image, seed, metadata_text
-
-# Define Gradio interface
-def interface_fn(prompt, negative_prompt, use_defaults, resolution, guidance_scale, num_inference_steps, seed, randomize_seed, progress=gr.Progress()):
-    image, seed, metadata_text = generate_image(prompt, negative_prompt, use_defaults, resolution, guidance_scale, num_inference_steps, seed, randomize_seed, progress)
-    return image, seed, gr.update(value=metadata_text)
-
-def reset_inputs():
-    return gr.update(value=''), gr.update(value=''), gr.update(value=True), gr.update(value='832x1216'), gr.update(value=7), gr.update(value=28), gr.update(value=0), gr.update(value=True), gr.update(value='')
-
-with gr.Blocks(title="UrangDiffusion 2.0 Demo", theme="NoCrypt/miku@1.2.1") as demo:
-    gr.HTML(
-        "<h1>UrangDiffusion 2.0 Demo</h1>"
-        "This demo is intended to showcase what the model is capable of and is not intended to be the main generation platform. Results produced with Diffusers are not the best, and it's highly recommended for you to get the model running inside Stable Diffusion WebUI or ComfyUI."
-        )
-    with gr.Row():
-        with gr.Column():
-            prompt_input = gr.Textbox(lines=2, placeholder="Enter prompt here", label="Prompt")
-            negative_prompt_input = gr.Textbox(lines=2, placeholder="Enter negative prompt here", label="Negative Prompt")
-            use_defaults_input = gr.Checkbox(label="Use Default Quality Tags and Negative Prompt", value=True)
-            resolution_input = gr.Radio(
-                choices=[
-                    "1024x1024", "1152x896", "896x1152", "1216x832", "832x1216",
-                    "1344x768", "768x1344", "1536x640", "640x1536"
-                ],
-                label="Resolution",
-                value="832x1216"
-            )
-            guidance_scale_input = gr.Slider(minimum=1, maximum=20, step=0.5, label="Guidance Scale", value=7)
-            num_inference_steps_input = gr.Slider(minimum=1, maximum=100, step=1, label="Number of Inference Steps", value=28)
-            seed_input = gr.Slider(minimum=0, maximum=999999999, step=1, label="Seed", value=0, interactive=True)
-            randomize_seed_input = gr.Checkbox(label="Randomize Seed", value=True)
-            generate_button = gr.Button("Generate")
-            reset_button = gr.Button("Reset")
-
-        with gr.Column():
-            output_image = gr.Image(type="pil", label="Generated Image")
-            with gr.Accordion("Parameters", open=False):
-                gr.Markdown(
-                    """
-                    This parameter is compatible with Stable Diffusion WebUI's parameter importer.
-                    """
-                )
-                metadata_textbox = gr.Textbox(lines=6, label="Image Parameters", interactive=False, max_lines=6)
-            gr.Markdown(
-                """
-                ### Recommended prompt formatting:
-                `1girl/1boy, character name, from what series, everything else in any order, best quality, amazing quality, very aesthetic,`
-
-                **PS:** `best quality, amazing quality, very aesthetic,` is automatically added when "Use Default Quality Tags and Negative Prompt" is enabled
-
-                ### Recommended settings:
-                - Steps: 25-30
-                - CFG: 5-7
-                """
-            )
-
-    generate_button.click(
-        interface_fn,
-        inputs=[
-            prompt_input, negative_prompt_input, use_defaults_input, resolution_input, guidance_scale_input, num_inference_steps_input, seed_input, randomize_seed_input
-        ],
-        outputs=[output_image, seed_input, metadata_textbox]
-    )
-    
-    reset_button.click(
-        reset_inputs,
-        inputs=[],
-        outputs=[
-            prompt_input, negative_prompt_input, use_defaults_input, resolution_input, guidance_scale_input, num_inference_steps_input, seed_input, randomize_seed_input, metadata_textbox
-        ]
-    )
-
-demo.queue(max_size=20).launch(share=False)
+import spaces
+import gradio as gr
+from hft2is import load_models, warm_models, generate_image_hf
+
+MODELS = ["kayfahaarukku/UrangDiffusion-2.0", "kayfahaarukku/irAsu-1.0"]
+load_models(MODELS)
+warm_models(MODELS)
+
+# Define Gradio interface
+def interface_fn(prompt, negative_prompt="", use_defaults=True, resolution="832x1216", guidance_scale=7, num_inference_steps=28, seed=0, randomize_seed=True, model_name=MODELS[0], progress=gr.Progress()):
+    image, seed, metadata_text = generate_image_hf(model_name, prompt, negative_prompt, use_defaults, resolution, guidance_scale, num_inference_steps, seed, randomize_seed, progress)
+    return image, seed, gr.update(value=metadata_text)
+
+def reset_inputs():
+    return gr.update(value=''), gr.update(value=''), gr.update(value=True), gr.update(value='832x1216'), gr.update(value=7), gr.update(value=28), gr.update(value=0), gr.update(value=True), gr.update(value=''), gr.update(value=MODELS[0])
+
+with gr.Blocks(title="UrangDiffusion 2.0 Demo", theme="NoCrypt/miku@1.2.1") as demo:
+    gr.HTML(
+        "<h1>UrangDiffusion 2.0 Demo</h1>"
+        "This demo is intended to showcase what the model is capable of and is not intended to be the main generation platform. Results produced with Diffusers are not the best, and it's highly recommended for you to get the model running inside Stable Diffusion WebUI or ComfyUI."
+        )
+    with gr.Row():
+        with gr.Column():
+            prompt_input = gr.Textbox(lines=2, placeholder="Enter prompt here", label="Prompt")
+            negative_prompt_input = gr.Textbox(lines=2, placeholder="Enter negative prompt here", label="Negative Prompt")
+            use_defaults_input = gr.Checkbox(label="Use Default Quality Tags and Negative Prompt", value=True)
+            resolution_input = gr.Radio(
+                choices=[
+                    "1024x1024", "1152x896", "896x1152", "1216x832", "832x1216",
+                    "1344x768", "768x1344", "1536x640", "640x1536"
+                ],
+                label="Resolution",
+                value="832x1216"
+            )
+            guidance_scale_input = gr.Slider(minimum=1, maximum=20, step=0.5, label="Guidance Scale", value=7)
+            num_inference_steps_input = gr.Slider(minimum=1, maximum=100, step=1, label="Number of Inference Steps", value=28)
+            seed_input = gr.Slider(minimum=0, maximum=999999999, step=1, label="Seed", value=0, interactive=True)
+            randomize_seed_input = gr.Checkbox(label="Randomize Seed", value=True)
+            model_name = gr.Dropdown(label="Model", choices=MODELS, value=MODELS[0])
+            generate_button = gr.Button("Generate")
+            reset_button = gr.Button("Reset")
+
+        with gr.Column():
+            output_image = gr.Image(type="pil", label="Generated Image")
+            with gr.Accordion("Parameters", open=False):
+                gr.Markdown(
+                    """
+                    This parameter is compatible with Stable Diffusion WebUI's parameter importer.
+                    """
+                )
+                metadata_textbox = gr.Textbox(lines=6, label="Image Parameters", interactive=False, max_lines=6)
+            gr.Markdown(
+                """
+                ### Recommended prompt formatting:
+                `1girl/1boy, character name, from what series, everything else in any order, best quality, amazing quality, very aesthetic,`
+
+                **PS:** `best quality, amazing quality, very aesthetic,` is automatically added when "Use Default Quality Tags and Negative Prompt" is enabled
+
+                ### Recommended settings:
+                - Steps: 25-30
+                - CFG: 5-7
+                """
+            )
+            examples = gr.Examples(
+                examples = [
+                    ["1girl"],
+                    ["beautiful sunset"],
+                ],
+                inputs=[prompt_input],
+                cache_examples=False,
+            )
+
+    generate_button.click(
+        interface_fn,
+        inputs=[
+            prompt_input, negative_prompt_input, use_defaults_input, resolution_input, guidance_scale_input, num_inference_steps_input, seed_input, randomize_seed_input, model_name,
+        ],
+        outputs=[output_image, seed_input, metadata_textbox]
+    )
+    
+    reset_button.click(
+        reset_inputs,
+        inputs=[],
+        outputs=[
+            prompt_input, negative_prompt_input, use_defaults_input, resolution_input, guidance_scale_input, num_inference_steps_input, seed_input, randomize_seed_input, metadata_textbox, model_name,
+        ]
+    )
+
+demo.launch(ssr_mode=False)

hft2is.py

@@ -0,0 +1,232 @@
+import gradio as gr
+import asyncio
+from threading import RLock
+from pathlib import Path
+import os
+from typing import Union
+
+
+HF_TOKEN = os.getenv("HF_TOKEN", None)
+server_timeout = 600
+inference_timeout = 600
+
+
+lock = RLock()
+loaded_models = {}
+
+
+def rename_image(image_path: Union[str, None], model_name: str, save_path: Union[str, None] = None):
+    import shutil
+    from datetime import datetime, timezone, timedelta
+    if image_path is None: return None
+    dt_now = datetime.now(timezone(timedelta(hours=9)))
+    filename = f"{model_name.split('/')[-1]}_{dt_now.strftime('%Y%m%d_%H%M%S')}.png"
+    try:
+        if Path(image_path).exists():
+            png_path = "image.png"
+            if str(Path(image_path).resolve()) != str(Path(png_path).resolve()): shutil.copy(image_path, png_path)
+            if save_path is not None:
+                new_path = str(Path(png_path).resolve().rename(Path(save_path).resolve()))
+            else:
+                new_path = str(Path(png_path).resolve().rename(Path(filename).resolve()))
+            return new_path
+        else:
+            return None
+    except Exception as e:
+        print(e)
+        return None
+
+
+# https://github.com/gradio-app/gradio/blob/main/gradio/external.py
+# https://huggingface.co/docs/huggingface_hub/package_reference/inference_client
+from typing import Literal
+def load_from_model(model_name: str, hf_token: Union[str, Literal[False], None] = None):
+    import httpx
+    import huggingface_hub
+    from gradio.exceptions import ModelNotFoundError, TooManyRequestsError
+    model_url = f"https://huggingface.co/{model_name}"
+    api_url = f"https://api-inference.huggingface.co/models/{model_name}"
+    print(f"Fetching model from: {model_url}")
+
+    headers = ({} if hf_token in [False, None] else {"Authorization": f"Bearer {hf_token}"})
+    response = httpx.request("GET", api_url, headers=headers)
+    if response.status_code != 200:
+        raise ModelNotFoundError(
+            f"Could not find model: {model_name}. If it is a private or gated model, please provide your Hugging Face access token (https://huggingface.co/settings/tokens) as the argument for the `hf_token` parameter."
+        )
+    p = response.json().get("pipeline_tag")
+    if p != "text-to-image": raise ModelNotFoundError(f"This model isn't for text-to-image or unsupported: {model_name}.")
+    headers["X-Wait-For-Model"] = "true"
+    kwargs = {}
+    if hf_token is not None: kwargs["token"] = hf_token
+    client = huggingface_hub.InferenceClient(model=model_name, headers=headers, timeout=server_timeout, **kwargs)
+    inputs = gr.components.Textbox(label="Input")
+    outputs = gr.components.Image(label="Output")
+    fn = client.text_to_image
+
+    def query_huggingface_inference_endpoints(*data, **kwargs):
+        try:
+            data = fn(*data, **kwargs)  # type: ignore
+        except huggingface_hub.utils.HfHubHTTPError as e:
+            print(e)
+            if "429" in str(e): raise TooManyRequestsError() from e
+        except Exception as e:
+            print(e)
+            raise Exception() from e
+        return data
+
+    interface_info = {
+        "fn": query_huggingface_inference_endpoints,
+        "inputs": inputs,
+        "outputs": outputs,
+        "title": model_name,
+    }
+    return gr.Interface(**interface_info)
+
+
+def load_model(model_name: str):
+    global loaded_models
+    global model_info_dict
+    if model_name in loaded_models.keys(): return loaded_models[model_name]
+    try:
+        loaded_models[model_name] = load_from_model(model_name, hf_token=HF_TOKEN)
+        print(f"Loaded: {model_name}")
+    except Exception as e:
+        if model_name in loaded_models.keys(): del loaded_models[model_name]
+        print(f"Failed to load: {model_name}")
+        print(e)
+        return None
+    return loaded_models[model_name]
+
+
+def load_models(models: list):
+    for model in models:
+        load_model(model)
+
+
+def warm_model(model_name: str):
+    model = load_model(model_name)
+    if model:
+        try:
+            print(f"Warming model: {model_name}")
+            infer_body(model, model_name, " ")
+        except Exception as e:
+            print(e)
+
+
+def warm_models(models: list[str]):
+    for model in models:
+        asyncio.new_event_loop().run_in_executor(None, warm_model, model)
+
+
+# https://huggingface.co/docs/api-inference/detailed_parameters
+# https://huggingface.co/docs/huggingface_hub/package_reference/inference_client
+def infer_body(client: Union[gr.Interface, object], model_str: str, prompt: str, neg_prompt: str = "",
+               height: int = 0, width: int = 0, steps: int = 0, cfg: int = 0, seed: int = -1):
+    png_path = "image.png"
+    kwargs = {}
+    if height > 0: kwargs["height"] = height
+    if width > 0: kwargs["width"] = width
+    if steps > 0: kwargs["num_inference_steps"] = steps
+    if cfg > 0: cfg = kwargs["guidance_scale"] = cfg
+    if seed == -1: kwargs["seed"] = randomize_seed()
+    else: kwargs["seed"] = seed
+    try:
+        if isinstance(client, gr.Interface): image = client.fn(prompt=prompt, negative_prompt=neg_prompt, **kwargs)
+        else: return None
+        if isinstance(image, tuple): return None
+        return save_image(image, png_path, model_str, prompt, neg_prompt, height, width, steps, cfg, seed)
+    except Exception as e:
+        print(e)
+        raise Exception(e) from e
+
+
+async def infer(model_name: str, prompt: str, neg_prompt: str ="", height: int = 0, width: int = 0,
+               steps: int = 0, cfg: int = 0, seed: int = -1,
+               save_path: str | None = None, timeout: float = inference_timeout):
+    model = load_model(model_name)
+    if not model: return None
+    task = asyncio.create_task(asyncio.to_thread(infer_body, model, model_name, prompt, neg_prompt,
+                                                 height, width, steps, cfg, seed))
+    await asyncio.sleep(0)
+    try:
+        result = await asyncio.wait_for(task, timeout=timeout)
+    except asyncio.TimeoutError as e:
+        print(e)
+        print(f"Task timed out: {model_name}")
+        if not task.done(): task.cancel()
+        result = None
+        raise Exception(f"Task timed out: {model_name}") from e
+    except Exception as e:
+        print(e)
+        if not task.done(): task.cancel()
+        result = None
+        raise Exception(e) from e
+    if task.done() and result is not None:
+        with lock:
+            image = rename_image(result, model_name, save_path)
+        return image
+    return None
+
+
+def save_image(image, savefile, modelname, prompt, nprompt, height=0, width=0, steps=0, cfg=0, seed=-1):
+    from PIL import Image, PngImagePlugin
+    import json
+    try:
+        metadata = {"prompt": prompt, "negative_prompt": nprompt, "Model": {"Model": modelname.split("/")[-1]}}
+        if steps > 0: metadata["num_inference_steps"] = steps
+        if cfg > 0: metadata["guidance_scale"] = cfg
+        if seed != -1: metadata["seed"] = seed
+        if width > 0 and height > 0: metadata["resolution"] = f"{width} x {height}"
+        metadata_str = json.dumps(metadata)
+        info = PngImagePlugin.PngInfo()
+        info.add_text("metadata", metadata_str)
+        image.save(savefile, "PNG", pnginfo=info)
+        return str(Path(savefile).resolve())
+    except Exception as e:
+        print(f"Failed to save image file: {e}")
+        raise Exception(f"Failed to save image file:") from e
+
+
+def randomize_seed():
+    from random import seed, randint
+    MAX_SEED = 2**32-1
+    seed()
+    rseed = randint(0, MAX_SEED)
+    return rseed
+
+
+def gen_image(model_name: str, prompt: str, neg_prompt: str = "", height: int = 0, width: int = 0,
+              steps: int = 0, cfg: int = 0, seed: int = -1):
+    if model_name in ["NA", ""]: return gr.update()
+    try:
+        loop = asyncio.get_running_loop()
+    except Exception:
+        loop = asyncio.new_event_loop()
+    try:
+        result = loop.run_until_complete(infer(model_name, prompt, neg_prompt, height, width,
+                                               steps, cfg, seed, None, inference_timeout))
+    except (Exception, asyncio.CancelledError) as e:
+        print(e)
+        print(f"Task aborted: {model_name}, Error: {e}")
+        result = None
+        raise gr.Error(f"Task aborted: {model_name}, Error: {e}")
+    finally:
+        loop.close()
+    return result
+
+
+def generate_image_hf(model_name: str, prompt: str, negative_prompt: str, use_defaults: bool, resolution: str,
+                      guidance_scale: float, num_inference_steps: int, seed: int, randomize_seed: bool, progress=gr.Progress()):
+
+    if randomize_seed: seed = -1
+    if use_defaults:
+        prompt = f"{prompt}, best quality, amazing quality, very aesthetic"
+        negative_prompt = f"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], {negative_prompt}"
+    
+    width, height = map(int, resolution.split('x'))
+    image = gen_image(model_name, prompt, negative_prompt, height, width, num_inference_steps, guidance_scale)
+
+    metadata_text = f"{prompt}\nNegative prompt: {negative_prompt}\nSteps: {num_inference_steps}, Sampler: Euler a, Size: {width}x{height}, Seed: {seed}, CFG scale: {guidance_scale}"
+
+    return image, seed, metadata_text

requirements.txt

@@ -1,12 +1 @@
-torch==2.0.1
-diffusers>=0.26.3
-transformers==4.38.1
-accelerate==0.27.2
-safetensors>=0.3.1
-gradio==4.20.0
-tqdm==4.66.4
-invisible-watermark==0.2.0
-Pillow==10.2.0
-spaces==0.24.0
-omegaconf==2.3.0
-timm==0.9.10
+huggingface_hub