Update app.py

app.py

@@ -1,166 +1,113 @@
 #!/usr/bin/env python
-#Patch0.1x
+# Patch3.09
 import os
 import random
 import uuid
-import json
+import gdown
 
 import gradio as gr
 import numpy as np
 from PIL import Image
 import spaces
 import torch
-from diffusers import DiffusionPipeline
-from typing import Tuple
+from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 
-# Check for the Model Base..//
-bad_words = json.loads(os.getenv('BAD_WORDS', "[]"))
-bad_words_negative = json.loads(os.getenv('BAD_WORDS_NEGATIVE', "[]"))
-default_negative = os.getenv("default_negative","")
+DESCRIPTION = """ """
 
-def check_text(prompt, negative=""):
-    for i in bad_words:
-        if i in prompt:
-            return True
-    for i in bad_words_negative:
-        if i in negative:
-            return True
-    return False
-# End of the - Prompt Con
-
-style_list = [
-    {
-        "name": "3840 x 2160",
-        "prompt": "hyper-realistic 8K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
-        "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
-    },
-    {
-        "name": "2560 x 1440",
-        "prompt": "hyper-realistic 4K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
-        "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
-    },
-    {
-        "name": "3D Model",
-        "prompt": "professional 3d model {prompt}. octane render, highly detailed, volumetric, dramatic lighting",
-        "negative_prompt": "ugly, deformed, noisy, low poly, blurry, painting",
-    },
-]
-
-styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
-STYLE_NAMES = list(styles.keys())
-DEFAULT_STYLE_NAME = "3840 x 2160"
-
-def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
-    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
-    if not negative:
-        negative = ""
-    return p.replace("{prompt}", positive), n + negative
-
-DESCRIPTION = """## MidJourney 3D
-"""
+def save_image(img):
+    unique_name = str(uuid.uuid4()) + ".png"
+    img.save(unique_name)
+    return unique_name
 
-if not torch.cuda.is_available():
-    DESCRIPTION += "\n<p>⚠️Running on CPU, This may not work on CPU.</p>"
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
 
 MAX_SEED = np.iinfo(np.int32).max
-CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "0") == "1"
-MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "2048"))
-USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
-ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
 
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+if not torch.cuda.is_available():
+    DESCRIPTION += "\n<p>Running on CPU, This Space may not work on CPU.</p>"
 
-NUM_IMAGES_PER_PROMPT = 1
+USE_TORCH_COMPILE = 0
+ENABLE_CPU_OFFLOAD = 0
 
 if torch.cuda.is_available():
-    pipe = DiffusionPipeline.from_pretrained(
-        "yodayo-ai/kivotos-xl-2.0",
+    pipe = StableDiffusionXLPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-xl-base-1.0",
         torch_dtype=torch.float16,
         use_safetensors=True,
-        add_watermarker=False,
-        variant="fp16"
     )
-    if ENABLE_CPU_OFFLOAD:
-        pipe.enable_model_cpu_offload()
-    else:
-        pipe.to(device)    
-        print("Loaded on Device!")
+    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+
+    # Download the LoRA weights from Google Drive
+    drive_folder_url = "https://drive.google.com/drive/folders/1ExL5VNChyYWXho1QbgNbOkTK3xc8mhHW"
+    weight_file_id = "18n6gF7Jda92MpqK7cYs0Gv2IqLAVnltZ"
+    weight_file_name = "pytorch_lora_weights.safetensors"
+    # Use gdown to download the file
+    gdown.download(f"https://drive.google.com/uc?id={weight_file_id}", weight_file_name, quiet=False)
     
-    if USE_TORCH_COMPILE:
-        pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
-        print("Model Compiled!")
-
-def save_image(img):
-    unique_name = str(uuid.uuid4()) + ".png"
-    img.save(unique_name)
-    return unique_name
+    pipe.load_lora_weights(weight_file_name, adapter_name="icon")
+    pipe.set_adapters("icon")
 
-def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    return seed
+    pipe.to("cuda")
 
 @spaces.GPU(enable_queue=True)
 def generate(
     prompt: str,
     negative_prompt: str = "",
     use_negative_prompt: bool = False,
-    style: str = DEFAULT_STYLE_NAME,
     seed: int = 0,
     width: int = 1024,
     height: int = 1024,
     guidance_scale: float = 3,
     randomize_seed: bool = False,
-    use_resolution_binning: bool = True,
     progress=gr.Progress(track_tqdm=True),
 ):
-    if check_text(prompt, negative_prompt):
-        raise ValueError("Prompt contains restricted words.")
-    
-    prompt, negative_prompt = apply_style(style, prompt, negative_prompt)
     seed = int(randomize_seed_fn(seed, randomize_seed))
-    generator = torch.Generator().manual_seed(seed)
 
     if not use_negative_prompt:
         negative_prompt = ""  # type: ignore
-    negative_prompt += default_negative    
-
-    options = {
-        "prompt": prompt,
-        "negative_prompt": negative_prompt,
-        "width": width,
-        "height": height,
-        "guidance_scale": guidance_scale,
-        "num_inference_steps": 25,
-        "generator": generator,
-        "num_images_per_prompt": NUM_IMAGES_PER_PROMPT,
-        "use_resolution_binning": use_resolution_binning,
-        "output_type": "pil",
-    }
-    
-    images = pipe(**options).images
 
+    images = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        width=width,
+        height=height,
+        guidance_scale=guidance_scale,
+        num_inference_steps=25,
+        num_images_per_prompt=1,
+        cross_attention_kwargs={"scale": 0.65},
+        output_type="pil",
+    ).images
     image_paths = [save_image(img) for img in images]
+    print(image_paths)
     return image_paths, seed
 
 examples = [
-    "A closeup of a cat, a window, in a rustic cabin, close up, with a shallow depth of field, with a vintage film grain, in the style of Annie Leibovitz and in the style of Wes Anderson. --ar 85:128 --v 6.0 --style raw",
-    "Daria Morgendorffer the main character of the animated series Daria, serious expression, very excites sultry look, so hot girl, beautiful charismatic girl, so hot shot, a woman wearing eye glasses, gorgeous figure, interesting shapes, life-size figures",
-    "Dark green large leaves of anthurium, close up, photography, aerial view, in the style of unsplash, hasselblad h6d400c  --ar 85:128 --v 6.0 --style raw",
-    "Closeup of blonde woman depth of field, bokeh, shallow focus, minimalism, fujifilm xh2s with Canon EF lens, cinematic --ar 85:128 --v 6.0 --style raw"
+    "1boy, male focus, sky, star (sky), night, pointing up, night sky, hood down, starry sky, hood, blue theme, outdoors, long sleeves, shooting star, hoodie, short hair, jacket, scenery, cloud, from behind, blue eyes, best quality, amazing quality, best aesthetic, absurdres",
+    "1boy, male focus, bishounen, holding sword, holding weapon, katana, sword, japanese clothes, haori, east asian architecture, solo, looking at viewer, expressionless, blue hair, purple eyes, long hair, best quality, amazing quality, best aesthetic, absurdres",
+    "1boy, male focus, holding drink, holding, drink, toned male, toned, pectorals, jacket, open jacket, open clothes, tank top, chain necklace, necklace, stud earrings, earrings, jewelry, cafe, plant, indoors, lens flare, solo, looking at viewer, open mouth, fang, white hair, yellow eyes, short hair, best quality, amazing quality, best aesthetic, absurdres, year 2023",
+    "1boy, male focus, dark-skinned male, dark skin, squatting, heart hands, bara, wooden floor, floor, indoors, gym uniform, sneakers, shoes, solo, looking at viewer, frown, sweatdrop, very short hair, best quality, amazing quality, best aesthetic, absurdres, year 2023",
+    "1boy, male focus, short hair, blue hair, blue eyes, graphic t-shirt, punk t-shirt, digital illustration, cyan and black, looking at viewer,  busy city street, belt, black pants, atmospheric lighting, midriff peek, night, blurry, best quality, amazing quality, best aesthetic, absurdres",
+    "Ultra realistic close up portrait ((beautiful pale cyberpunk female with heavy black eyeliner)), blue eyes, shaved side haircut, hyper detail, cinematic lighting, magic neon, dark red city, Canon EOS R3, nikon, f/1.4, ISO 200, 1/160s, 8K, RAW, unedited, symmetrical balance, in-frame, 8K"
 ]
 
 css = '''
-.gradio-container{max-width: 580px !important}
+.gradio-container{max-width: 600px !important}
 h1{text-align:center}
+footer {
+    visibility: hidden
+}
 '''
-with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
+with gr.Blocks(css=css, theme="ParityError/Anime") as demo:
     gr.Markdown(DESCRIPTION)
     gr.DuplicateButton(
         value="Duplicate Space for private use",
         elem_id="duplicate-button",
-        visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
+        visible=False,
     )
+
     with gr.Group():
         with gr.Row():
             prompt = gr.Text(
@@ -170,33 +117,18 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
                 placeholder="Enter your prompt",
                 container=False,
             )
-            run_button = gr.Button("Run")
-        result = gr.Gallery(label="Result", columns=1, preview=True)
+            run_button = gr.Button("Run", scale=0)
+        result = gr.Gallery(label="Result", columns=1, preview=True, show_label=False)
     with gr.Accordion("Advanced options", open=False):
-        use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True, visible=True)
+        use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True)
         negative_prompt = gr.Text(
             label="Negative prompt",
-            max_lines=1,
+            lines=4,
+            max_lines=6,
+            value="""(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation""",
             placeholder="Enter a negative prompt",
-            value="deformed iris, deformed pupils, semi-realistic, text, close up, cropped, out of frame, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck",
             visible=True,
         )
-        with gr.Row():
-            num_inference_steps = gr.Slider(
-                label="Steps",
-                minimum=10,
-                maximum=60,
-                step=1,
-                value=30,
-            )
-        with gr.Row():
-            num_images_per_prompt = gr.Slider(
-                label="Images",
-                minimum=1,
-                maximum=5,
-                step=1,
-                value=2,
-            )
         seed = gr.Slider(
             label="Seed",
             minimum=0,
@@ -229,21 +161,13 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
                 step=0.1,
                 value=6,
             )
-    with gr.Row(visible=True):
-        style_selection = gr.Radio(
-            show_label=True,
-            container=True,
-            interactive=True,
-            choices=STYLE_NAMES,
-            value=DEFAULT_STYLE_NAME,
-            label="Image Style",
-        )
+
     gr.Examples(
         examples=examples,
         inputs=prompt,
         outputs=[result, seed],
         fn=generate,
-        cache_examples=CACHE_EXAMPLES,
+        cache_examples=False,
     )
 
     use_negative_prompt.change(
@@ -264,7 +188,6 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             prompt,
             negative_prompt,
             use_negative_prompt,
-            style_selection,
             seed,
             width,
             height,
@@ -274,6 +197,6 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
         outputs=[result, seed],
         api_name="run",
     )
-
+    
 if __name__ == "__main__":
-    demo.queue(max_size=50).launch()
+    demo.queue(max_size=20).launch(show_api=False, debug=False)