Update app.py

app.py

@@ -4,37 +4,27 @@ import numpy as np
 import modin.pandas as pd
 from PIL import Image
 from huggingface_hub import hf_hub_download
-from diffusers import StableCascadeDecoderPipeline, StableCascadePriorPipeline
+from diffusers import StableDiffusion3Pipeline
 
 device = 'cuda' #if torch.cuda.is_available() else 'cpu'
 torch.cuda.max_memory_allocated(device=device)
 torch.cuda.empty_cache()
-prior = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", variant="bf16", torch_dtype=torch.bfloat16)
-prior.enable_xformers_memory_efficient_attention()
-prior = prior.to(device)
-decoder = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade", variant="bf16", torch_dtype=torch.float16)
-decoder.enable_xformers_memory_efficient_attention()
-decoder.enable_model_cpu_offload()
-torch.cuda.empty_cache()
+SD3 = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16)
+SD3.enable_xformers_memory_efficient_attention()
+SD3 = SD3.to(device)
+
                     
 def genie (Prompt, negative_prompt, height, width, scale, steps, d_steps, seed):
     generator = np.random.seed(0) if seed == 0 else torch.manual_seed(seed)
     torch.cuda.empty_cache()
-    prior_image=prior(
+    image=SD3(
     prompt=Prompt,
     height=height,
     width=width,
     negative_prompt=negative_prompt,
     guidance_scale=scale,
     num_images_per_prompt=1,
-    num_inference_steps=steps)
-    image=decoder(
-    image_embeddings=prior_image.image_embeddings.to(torch.float16),
-    prompt=Prompt,
-    negative_prompt=negative_prompt,
-    guidance_scale=0.0,
-    output_type="pil",
-    num_inference_steps=d_steps).images[0]
+    num_inference_steps=steps).images[0]
     return image
     
 gr.Interface(fn=genie, inputs=[gr.Textbox(label='What you want the AI to generate. 77 Token Limit.'),