Update app.py

app.py

@@ -1,4 +1,5 @@
 from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler, AutoencoderKL
+from transformers import CLIPTextModel, CLIPTokenizer
 import torch
 import gradio as gr
 import spaces
@@ -18,9 +19,23 @@ def generate_image(prompt, negative_prompt, num_inference_steps=30, guidance_sca
         
     else:
         model_id = "SG161222/Realistic_Vision_V6.0_B1_noVAE"
+
+    text_encoder = CLIPTextModel.from_pretrained(
+        model_id,
+        subfolder="text_encoder"
+    ).to("cuda")
+    
+    tokenizer = CLIPTokenizer.from_pretrained(
+        model_id,
+        subfolder="tokenizer"
+    )
         
-        
-    pipe = DiffusionPipeline.from_pretrained(model_id, vae=vae).to("cuda")
+    pipe = DiffusionPipeline.from_pretrained(
+        model_id,
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        vae=vae
+    ).to("cuda")
     
     if model == "Real6.0":
         pipe.safety_checker = lambda images, **kwargs: (images, [False] * len(images))
@@ -33,16 +48,35 @@ def generate_image(prompt, negative_prompt, num_inference_steps=30, guidance_sca
         use_karras_sigmas=True
     )
 
+    text_inputs = tokenizer(
+        prompt,
+        padding="max_length",
+        max_length=tokenizer.model_max_length,
+        truncation=True,
+        return_tensors="pt"
+    ).to("cuda")
+    
+    negative_text_inputs = tokenizer(
+        negative_prompt,
+        padding="max_length",
+        max_length=tokenizer.model_max_length,
+        truncation=True,
+        return_tensors="pt"
+    ).to("cuda")
+
+    prompt_embeds = text_encoder(text_inputs.input_ids)[0]
+    negative_prompt_embeds = text_encoder(negative_text_inputs.input_ids)[0]
+
     
     # Generate the image
     result = pipe(
-        prompt = prompt,
-        negative_prompt = negative_prompt,
-        cross_attention_kwargs = {"scale":1},
-        num_inference_steps = num_inference_steps,
-        guidance_scale = guidance_scale,
-        width = width,
-        height = height,
+        prompt_embeds=prompt_embeds,
+        negative_prompt_embeds=negative_prompt_embeds,
+        cross_attention_kwargs={"scale": 1},
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        width=width,
+        height=height,
         num_images_per_prompt=num_images
     )