Update app.py

app.py

@@ -12,9 +12,8 @@ import gradio as gr
 # Download the model files
 ckpt_dir = snapshot_download(repo_id="Kwai-Kolors/Kolors")
 
-# Function to load models
+# Function to load models on demand
 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()
@@ -23,17 +22,16 @@ def load_models():
     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")
-
-    return pipe
+    return StableDiffusionXLPipeline(
+        vae=vae,
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        unet=unet,
+        scheduler=scheduler,
+        force_zeros_for_empty_prompt=False
+    ).to("cuda")
 
+# Create a global variable to hold the pipeline
 pipe = load_models()
 
 @spaces.GPU(duration=200)
@@ -43,9 +41,10 @@ def generate_image(prompt, negative_prompt, height, width, num_inference_steps,
     else:
         seed = int(seed)  # Ensure seed is an integer
     
-    # Move the model to the GPU for inference
+    # Move the model to the GPU for inference and clear unnecessary variables
     with torch.no_grad():
-        image = pipe(
+        generator = torch.Generator(pipe.device).manual_seed(seed)
+        result = pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,
             height=height,
@@ -53,20 +52,13 @@ def generate_image(prompt, negative_prompt, height, width, num_inference_steps,
             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
-
+            generator=generator
+        )
+        image = result.images
+    
     return image, seed
 
-description = """
-<p align="center">Effective Training of Diffusion Model for Photorealistic Text-to-Image Synthesis</p>
-<p><center>
-<a href="https://kolors.kuaishou.com/" target="_blank">[Official Website]</a>
-<a href="https://github.com/Kwai-Kolors/Kolors/blob/master/imgs/Kolors_paper.pdf" target="_blank">[Tech Report]</a>
-<a href="https://huggingface.co/Kwai-Kolors/Kolors" target="_blank">[Model Page]</a>
-<a href="https://github.com/Kwai-Kolors/Kolors" target="_blank">[Github]</a>
-</center></p>
-"""
+
 
 # Gradio interface
 iface = gr.Interface(
@@ -90,8 +82,7 @@ iface = gr.Interface(
         gr.Number(label="Seed Used")
     ],
     title="Kolors",
-    description=description,
     theme='bethecloud/storj_theme',
 )
 
-iface.launch()  # Set debug=False for production
+iface.launch()