Update app_demo.py

app_demo.py

@@ -7,16 +7,13 @@ import gradio as gr
 import numpy as np
 import PIL.Image
 import torch
-#from diffusers import DiffusionPipeline
 from diffusers import StableDiffusionPipeline
-from tqdm import tqdm
-from safetensors.torch import load_file
 from concurrent.futures import ThreadPoolExecutor
 import uuid
-#import cv2
-model_id = "Lykon/dreamshaper-xl-v2-turbo" #"Lykon/dreamshaper-7" #"openskyml/lcm-lora-sdxl-turbo" #"SimianLuo/LCM_Dreamshaper_v7"
+
+model_id = "Lykon/dreamshaper-xl-v2-turbo"
 DESCRIPTION = '''# Fast Stable Diffusion CPU with Latent Consistency Model
-Distilled from [Dreamshaper v7](https://huggingface.co/Lykon/dreamshaper-7) fine-tune of [Stable Diffusion v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5) with only 4,000 training iterations (~32 A100 GPU Hours). [Project page](https://latent-consistency-models.github.io)
+Distilled from [Dreamshaper v7](https://huggingface.co/Lykon/dreamshaper-7) fine‑tune of SD v1-5.
 '''
 if not torch.cuda.is_available():
     DESCRIPTION += "\n<p>running on CPU.</p>"
@@ -24,44 +21,15 @@ if not torch.cuda.is_available():
 MAX_SEED = np.iinfo(np.int32).max
 CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES") == "1"
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "768"))
-USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE") == "1"
-DTYPE = torch.float32  # torch.float16 works as well, but pictures seem to be a bit worse
-
-#pipe = DiffusionPipeline.from_pretrained("SimianLuo/LCM_Dreamshaper_v7", custom_pipeline="latent_consistency_txt2img", custom_revision="main")
+DTYPE = torch.float32
 
-
-#"SimianLuo/LCM_Dreamshaper_v7"
-'''
-pipe = DiffusionPipeline.from_pretrained(  model_id  ,    custom_pipeline=model_id,
-    custom_revision="main",
-    low_cpu_mem_usage=True,
-    safety_checker=None,  # Disable NSFW filter
-    requires_safety_checker=False,  # Skip warning
+# Load pipeline once, disabling NSFW filter at construction time
+pipe = StableDiffusionPipeline.from_pretrained(
+    model_id,
+    safety_checker=None,
+    torch_dtype=DTYPE,
     use_safetensors=True
-)
-#pipe.to(torch_device="cpu",torch_dtype="float16",  torch_dtype=DTYPE)
-pipe.to(torch_dtype="float32" )
-pipe.to("cpu")
-'''
-
-#from diffusers import StableDiffusionPipeline
-
-
-
-pipe = StableDiffusionPipeline.from_pretrained(model_id,  safety_checker= None)
-prompt = "A futuristic cityscape at sunset"
-output = pipe(
-    prompt=prompt,
-    negative_prompt="",      # ← prevents added_cond_kwargs from being None
-    #num_inference_steps=50,  # adjust as you like
-    #guidance_scale=7.5       # ditto
-    requires_safety_checker=False
-    
-)
-image = output.images[0]
-#image = pipe(prompt).images[0]
-image.show()
-
+).to("cpu")
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
@@ -71,14 +39,14 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
 def save_image(img, profile: gr.OAuthProfile | None, metadata: dict):
     unique_name = str(uuid.uuid4()) + '.png'
     img.save(unique_name)
-    #gr_user_history.save_image(label=metadata["prompt"], image=img, profile=profile, metadata=metadata)
     return unique_name
 
 def save_images(image_array, profile: gr.OAuthProfile | None, metadata: dict):
-    paths = []
     with ThreadPoolExecutor() as executor:
-        paths = list(executor.map(save_image, image_array, [profile]*len(image_array), [metadata]*len(image_array)))
-    return paths
+        return list(executor.map(
+            lambda args: save_image(*args),
+            zip(image_array, [profile]*len(image_array), [metadata]*len(image_array))
+        ))
 
 def generate(
     prompt: str,
@@ -91,143 +59,83 @@ def generate(
     randomize_seed: bool = False,
     progress = gr.Progress(track_tqdm=True),
     profile: gr.OAuthProfile | None = None,
-) -> PIL.Image.Image:
+) -> tuple[list[str], int]:
+    # prepare seed
     seed = randomize_seed_fn(seed, randomize_seed)
     torch.manual_seed(seed)
+
     start_time = time.time()
-    result = pipe(
+    # **Call the pipeline with only supported kwargs:**
+    outputs = pipe(
         prompt=prompt,
-        width=width,
+        negative_prompt="",               # required to avoid NoneType in UNet
         height=height,
+        width=width,
         guidance_scale=guidance_scale,
-        negative_prompt="",      
-        safety_checker= None,    
-        requires_safety_checker=False,
         num_inference_steps=num_inference_steps,
         num_images_per_prompt=num_images,
-        lcm_origin_steps=50,
         output_type="pil",
     ).images
-    paths = save_images(result, profile, metadata={"prompt": prompt, "seed": seed, "width": width, "height": height, "guidance_scale": guidance_scale, "num_inference_steps": num_inference_steps})
-    print(time.time() - start_time)
+
+    latency = time.time() - start_time
+    print(f"Generation took {latency:.2f} seconds")
+
+    paths = save_images(
+        outputs,
+        profile,
+        metadata={
+            "prompt": prompt,
+            "seed": seed,
+            "width": width,
+            "height": height,
+            "guidance_scale": guidance_scale,
+            "num_inference_steps": num_inference_steps,
+        }
+    )
+
     return paths, seed
 
 examples = [
-    "portrait photo of a girl, photograph, highly detailed face, depth of field, moody light, golden hour, style by Dan Winters, Russell James, Steve McCurry, centered, extremely detailed, Nikon D850, award winning photography",
-    "Self-portrait oil painting, a beautiful cyborg with golden hair, 8k",
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "A photo of beautiful mountain with realistic sunset and blue lake, highly detailed, masterpiece",
+    "A futuristic cityscape at sunset",
+    "Steampunk airship over mountains",
+    "Portrait of a cyborg queen, hyper‑detailed",
 ]
 
 with gr.Blocks(css="style.css") as demo:
     gr.Markdown(DESCRIPTION)
-    gr.DuplicateButton(
-        value="Duplicate Space for private use",
-        elem_id="duplicate-button",
-        visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
-    )
-    gr.HTML(
-            f"""
-            <p id="project-links" align="center">
-                <a href='https://huggingface.co/spaces/charliebaby2023/Fast_Stable_diffusion_CPU/edit/main/app_demo.py'>Edit this py file</a>
-            </p>
-            """
-        )
     with gr.Group():
         with gr.Row():
             prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
                 placeholder="Enter your prompt",
+                show_label=False,
                 container=False,
             )
             run_button = gr.Button("Run", scale=0)
-        result = gr.Gallery(
-            label="Generated images", show_label=False, elem_id="gallery", grid=[2]
+        gallery = gr.Gallery(
+            label="Generated images",
+            show_label=False,
+            elem_id="gallery",
+            grid=[2]
         )
+
     with gr.Accordion("Advanced options", open=False):
-        seed = gr.Slider(
-            label="Seed",
-            minimum=0,
-            maximum=MAX_SEED,
-            step=1,
-            value=0,
-            randomize=True
-        )
+        seed = gr.Slider(0, MAX_SEED, value=0, step=1, randomize=True, label="Seed")
         randomize_seed = gr.Checkbox(label="Randomize seed across runs", value=True)
         with gr.Row():
-            width = gr.Slider(
-                label="Width",
-                minimum=256,
-                maximum=MAX_IMAGE_SIZE,
-                step=32,
-                value=512,
-            )
-            height = gr.Slider(
-                label="Height",
-                minimum=256,
-                maximum=MAX_IMAGE_SIZE,
-                step=32,
-                value=512,
-            )
+            width = gr.Slider(256, MAX_IMAGE_SIZE, value=512, step=32, label="Width")
+            height = gr.Slider(256, MAX_IMAGE_SIZE, value=512, step=32, label="Height")
         with gr.Row():
-            guidance_scale = gr.Slider(
-                label="Guidance scale for base",
-                minimum=2,
-                maximum=14,
-                step=0.1,
-                value=8.0,
-            )
-            num_inference_steps = gr.Slider(
-                label="Number of inference steps for base",
-                minimum=1,
-                maximum=8,
-                step=1,
-                value=4,
-            )
-        with gr.Row():
-            num_images = gr.Slider(
-                label="Number of images",
-                minimum=1,
-                maximum=8,
-                step=1,
-                value=1,
-                visible=True,
-            )
+            guidance_scale = gr.Slider(2.0, 14.0, value=8.0, step=0.1, label="Guidance Scale")
+            num_inference_steps = gr.Slider(1, 8, value=4, step=1, label="Inference Steps")
+        num_images = gr.Slider(1, 8, value=1, step=1, label="Number of Images")
 
-    with gr.Accordion("Past generations", open=False):
-        tr = gr.Textbox(value="ol")
-    
     gr.Examples(
         examples=examples,
         inputs=prompt,
-        outputs=result,
+        outputs=gallery,
         fn=generate,
         cache_examples=CACHE_EXAMPLES,
     )
 
-    gr.on(
-        triggers=[
-            prompt.submit,
-            run_button.click,
-        ],
-        fn=generate,
-        inputs=[
-            prompt,
-            seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-            num_images,
-            randomize_seed
-        ],
-        outputs=[result, seed],
-        api_name="run",
-    )
-
-if __name__ == "__main__":
-    demo.queue(api_open=False)
-    # demo.queue(max_size=20).launch()
+    demo.queue()
     demo.launch()