Update app.py

app.py

@@ -1,9 +1,7 @@
-#!/usr/bin/env python
-# Patch3.09
 import os
 import random
 import uuid
-import gdown
+import json
 
 import gradio as gr
 import numpy as np
@@ -12,7 +10,32 @@ import spaces
 import torch
 from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 
-DESCRIPTION = """ """
+# Use environment variables for flexibility
+MODEL_ID = os.getenv("MODEL_ID", "SG161222/RealVisXL_V4.0_Lightning")
+MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "4096"))
+USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
+ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
+BATCH_SIZE = int(os.getenv("BATCH_SIZE", "1"))  # Allow generating multiple images at once
+
+# Determine device and load model outside of function for efficiency
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+pipe = StableDiffusionXLPipeline.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    use_safetensors=True,
+    add_watermarker=False,
+).to(device)
+pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+
+# Torch compile for potential speedup (experimental)
+if USE_TORCH_COMPILE:
+    pipe.compile()
+
+# CPU offloading for larger RAM capacity (experimental)
+if ENABLE_CPU_OFFLOAD:
+    pipe.enable_model_cpu_offload()
+
+MAX_SEED = np.iinfo(np.int32).max
 
 def save_image(img):
     unique_name = str(uuid.uuid4()) + ".png"
@@ -24,90 +47,71 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
         seed = random.randint(0, MAX_SEED)
     return seed
 
-MAX_SEED = np.iinfo(np.int32).max
-
-if not torch.cuda.is_available():
-    DESCRIPTION += "\n<p>Running on CPU, This Space may not work on CPU.</p>"
-
-USE_TORCH_COMPILE = 0
-ENABLE_CPU_OFFLOAD = 0
-
-if torch.cuda.is_available():
-    pipe = StableDiffusionXLPipeline.from_pretrained(
-        "stabilityai/stable-diffusion-xl-base-1.0",
-        torch_dtype=torch.float16,
-        use_safetensors=True,
-    )
-    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)
-    
-    pipe.load_lora_weights(weight_file_name, adapter_name="icon")
-    pipe.set_adapters("icon")
-
-    pipe.to("cuda")
-
-@spaces.GPU(enable_queue=True)
+@spaces.GPU(duration=35, enable_queue=True)
 def generate(
     prompt: str,
     negative_prompt: str = "",
     use_negative_prompt: bool = False,
-    seed: int = 0,
+    seed: int = 1,
     width: int = 1024,
     height: int = 1024,
     guidance_scale: float = 3,
+    num_inference_steps: int = 30,
     randomize_seed: bool = False,
+    use_resolution_binning: bool = True, 
+    num_images: int = 1,  # Number of images to generate
     progress=gr.Progress(track_tqdm=True),
 ):
     seed = int(randomize_seed_fn(seed, randomize_seed))
+    generator = torch.Generator(device=device).manual_seed(seed)
+
+    # Improved options handling
+    options = {
+        "prompt": [prompt] * num_images,
+        "negative_prompt": [negative_prompt] * num_images if use_negative_prompt else None,
+        "width": width,
+        "height": height,
+        "guidance_scale": guidance_scale,
+        "num_inference_steps": num_inference_steps,
+        "generator": generator,
+        "output_type": "pil",
+    }
+
+    # Use resolution binning for faster generation with less VRAM usage
+    if use_resolution_binning:
+        options["use_resolution_binning"] = True
+
+    # Generate images potentially in batches
+    images = []
+    for i in range(0, num_images, BATCH_SIZE):
+        batch_options = options.copy()
+        batch_options["prompt"] = options["prompt"][i:i+BATCH_SIZE]
+        if "negative_prompt" in batch_options:
+            batch_options["negative_prompt"] = options["negative_prompt"][i:i+BATCH_SIZE]
+        images.extend(pipe(**batch_options).images)
 
-    if not use_negative_prompt:
-        negative_prompt = ""  # type: ignore
-
-    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 = [
-    "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"
+    "a cat eating a piece of cheese",
+    "a ROBOT riding a BLUE horse on Mars, photorealistic, 4k",
+    "Ironman VS Hulk, ultrarealistic",
+    "Astronaut in a jungle, cold color palette, oil pastel, detailed, 8k",
+    "An alien holding a sign board containing the word 'Flash', futuristic, neonpunk",
+    "Kids going to school, Anime style"
 ]
 
 css = '''
-.gradio-container{max-width: 600px !important}
+.gradio-container{max-width: 700px !important}
 h1{text-align:center}
 footer {
     visibility: hidden
 }
 '''
-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=False,
-    )
 
+with gr.Blocks(css=css) as demo:
+    gr.Markdown("""# SDXL Flash""")
     with gr.Group():
         with gr.Row():
             prompt = gr.Text(
@@ -118,56 +122,68 @@ with gr.Blocks(css=css, theme="ParityError/Anime") as demo:
                 container=False,
             )
             run_button = gr.Button("Run", scale=0)
-        result = gr.Gallery(label="Result", columns=1, preview=True, show_label=False)
+        result = gr.Gallery(label="Result", columns=1, show_label=False) 
     with gr.Accordion("Advanced options", open=False):
-        use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True)
-        negative_prompt = gr.Text(
-            label="Negative prompt",
-            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",
-            visible=True,
+        num_images = gr.Slider(
+            label="Number of Images",
+            minimum=1,
+            maximum=4,
+            step=1,
+            value=1,
         )
+        with gr.Row():
+            use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True)
+            negative_prompt = gr.Text(
+                label="Negative prompt",
+                max_lines=5,
+                lines=4,
+                placeholder="Enter a negative prompt",
+                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, NSFW",
+                visible=True,
+            )
         seed = gr.Slider(
             label="Seed",
             minimum=0,
             maximum=MAX_SEED,
             step=1,
             value=0,
-            visible=True
         )
         randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
         with gr.Row(visible=True):
             width = gr.Slider(
                 label="Width",
                 minimum=512,
-                maximum=2048,
-                step=8,
+                maximum=MAX_IMAGE_SIZE,
+                step=64,
                 value=1024,
             )
             height = gr.Slider(
                 label="Height",
                 minimum=512,
-                maximum=2048,
-                step=8,
+                maximum=MAX_IMAGE_SIZE,
+                step=64,
                 value=1024,
             )
         with gr.Row():
             guidance_scale = gr.Slider(
                 label="Guidance Scale",
                 minimum=0.1,
-                maximum=20.0,
+                maximum=6,
                 step=0.1,
-                value=6,
+                value=3.0,
+            )
+            num_inference_steps = gr.Slider(
+                label="Number of inference steps",
+                minimum=1,
+                maximum=15,
+                step=1,
+                value=8,
             )
 
     gr.Examples(
         examples=examples,
         inputs=prompt,
-        outputs=[result, seed],
-        fn=generate,
-        cache_examples=False,
+        cache_examples=False
     )
 
     use_negative_prompt.change(
@@ -192,11 +208,13 @@ with gr.Blocks(css=css, theme="ParityError/Anime") as demo:
             width,
             height,
             guidance_scale,
+            num_inference_steps,
             randomize_seed,
+            num_images
         ],
         outputs=[result, seed],
         api_name="run",
     )
-    
+
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch(show_api=False, debug=False)
+    demo.queue(max_size=50).launch()