Add application file

app.py

@@ -49,7 +49,8 @@ def load_model(model_name):
     else:
         current_model = StableDiffusion3Pipeline.from_pretrained(model_paths[model_name], torch_dtype=torch.bfloat16).to("cuda")
     
-    return current_model
+    return current_model.to('cuda')
+
 
 def generate_content(prompt, model_name, guidance_scale=7.5, num_inference_steps=50, use_cfg_zero_star=True, use_zero_init=True, zero_steps=0, seed=None, compare_mode=False):
     model = load_model(model_name)