Update app.py

app.py

@@ -4,18 +4,36 @@ import logging
 import torch
 from PIL import Image
 import spaces
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, AutoPipelineForText2Image
 from diffusers import StableDiffusion3Pipeline # pip install diffusers>=0.31.0
 import copy
 import random
 import time
+from huggingface_hub import login
+hf_token = os.environ.get("HF_TOKEN")
+login(token=hf_token)
+
+torch.set_float32_matmul_precision("high")
+
+torch._inductor.config.conv_1x1_as_mm = True
+torch._inductor.config.coordinate_descent_tuning = True
+torch._inductor.config.epilogue_fusion = False
+torch._inductor.config.coordinate_descent_check_all_directions = True
+
 
 # Load LoRAs from JSON file
 with open('loras.json', 'r') as f:
     loras = json.load(f)
-
+    
 # Initialize the base model
-pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3.5-large", torch_dtype=torch.bfloat16)
+base_model = "stabilityai/stable-diffusion-3.5-large"
+pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/stable-diffusion-3.5-large", torch_dtype=torch.bfloat16)
+
+pipe.transformer.to(memory_format=torch.channels_last)
+pipe.vae.to(memory_format=torch.channels_last)
+
+pipe.transformer = torch.compile(pipe.transformer, mode="max-autotune", fullgraph=True)
+pipe.vae.decode = torch.compile(pipe.vae.decode, mode="max-autotune", fullgraph=True)
 
 MAX_SEED = 2**32-1
 
@@ -71,7 +89,7 @@ def infer(prompt, negative_prompt, trigger_word, steps, seed, cfg_scale, width,
             width=width,
             height=height,
             generator=generator,
-            lora_scale=lora_scale,
+            joint_attention_kwargs={"scale": lora_scale},
         ).images[0]
     return image