Update app.py

app.py

@@ -3,21 +3,19 @@ import numpy as np
 import random
 import torch
 from diffusers import StableDiffusion3Pipeline, SD3Transformer2DModel, FlowMatchEulerDiscreteScheduler
+import spaces
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float16
 
-repo = "diffusers-internal-dev/private-model"
-transformer_repo= "diffusers-internal-dev/pm-revamp"
+repo = "stabilityai/stable-diffusion-3-medium"
 
-transformer = SD3Transformer2DModel.from_pretrained(transformer_repo, subfolder="transformer", torch_dtype=dtype)
-
-pipe = StableDiffusion3Pipeline.from_pretrained(repo, transformer=transformer, torch_dtype=torch.float16).to(device)
-pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(pipe.scheduler.config, shift=3.0)
+pipe = StableDiffusion3Pipeline.from_pretrained(repo, torch_dtype=torch.float16).to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1344
 
+@spaces.GPU
 def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress(track_tqdm=True)):
 
     if randomize_seed: