Update app.py

app.py

@@ -1,4 +1,3 @@
-import spaces
 import os
 import torch
 import random
@@ -13,23 +12,29 @@ import gradio as gr
 # Download the model files
 ckpt_dir = snapshot_download(repo_id="Kwai-Kolors/Kolors")
 
-# Load the models
-text_encoder = ChatGLMModel.from_pretrained(
-    os.path.join(ckpt_dir, 'text_encoder'),
-    torch_dtype=torch.float16).half()
-tokenizer = ChatGLMTokenizer.from_pretrained(os.path.join(ckpt_dir, 'text_encoder'))
-vae = AutoencoderKL.from_pretrained(os.path.join(ckpt_dir, "vae"), revision=None).half()
-scheduler = EulerDiscreteScheduler.from_pretrained(os.path.join(ckpt_dir, "scheduler"))
-unet = UNet2DConditionModel.from_pretrained(os.path.join(ckpt_dir, "unet"), revision=None).half()
+# Function to load models
+def load_models():
+    # Load models on demand to reduce initial memory footprint
+    text_encoder = ChatGLMModel.from_pretrained(
+        os.path.join(ckpt_dir, 'text_encoder'),
+        torch_dtype=torch.float16).half()
+    tokenizer = ChatGLMTokenizer.from_pretrained(os.path.join(ckpt_dir, 'text_encoder'))
+    vae = AutoencoderKL.from_pretrained(os.path.join(ckpt_dir, "vae"), revision=None).half()
+    scheduler = EulerDiscreteScheduler.from_pretrained(os.path.join(ckpt_dir, "scheduler"))
+    unet = UNet2DConditionModel.from_pretrained(os.path.join(ckpt_dir, "unet"), revision=None).half()
 
-pipe = StableDiffusionXLPipeline(
-        vae=vae,
-        text_encoder=text_encoder,
-        tokenizer=tokenizer,
-        unet=unet,
-        scheduler=scheduler,
-        force_zeros_for_empty_prompt=False)
-pipe = pipe.to("cuda")
+    pipe = StableDiffusionXLPipeline(
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            unet=unet,
+            scheduler=scheduler,
+            force_zeros_for_empty_prompt=False)
+    pipe = pipe.to("cuda")
+
+    return pipe
+
+pipe = load_models()
 
 @spaces.GPU(duration=200)
 def generate_image(prompt, negative_prompt, height, width, num_inference_steps, guidance_scale, num_images_per_prompt, use_random_seed, seed, progress=gr.Progress(track_tqdm=True)):
@@ -38,16 +43,19 @@ def generate_image(prompt, negative_prompt, height, width, num_inference_steps,
     else:
         seed = int(seed)  # Ensure seed is an integer
     
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        height=height,
-        width=width,
-        num_inference_steps=num_inference_steps,
-        guidance_scale=guidance_scale,
-        num_images_per_prompt=num_images_per_prompt,
-        generator=torch.Generator(pipe.device).manual_seed(seed)
-    ).images
+    # Move the model to the GPU for inference
+    with torch.no_grad():
+        image = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            num_images_per_prompt=num_images_per_prompt,
+            generator=torch.Generator(pipe.device).manual_seed(seed)
+        ).images
+
     return image, seed
 
 description = """
@@ -86,4 +94,4 @@ iface = gr.Interface(
     theme='bethecloud/storj_theme',
 )
 
-iface.launch(debug=True)
+iface.launch()  # Set debug=False for production