Update app.py

app.py

@@ -18,44 +18,12 @@ if not torch.cuda.is_available():
 
 MAX_SEED = np.iinfo(np.int32).max
 CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES") == "1"
-MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1024"))
+MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1824"))
 USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE") == "1"
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1"
 ENABLE_REFINER = os.getenv("ENABLE_REFINER", "1") == "1"
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-if torch.cuda.is_available():
-    vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
-    pipe = DiffusionPipeline.from_pretrained(
-        "stabilityai/stable-diffusion-xl-base-1.0",
-        vae=vae,
-        torch_dtype=torch.float16,
-        use_safetensors=True,
-        variant="fp16",
-    )
-    if ENABLE_REFINER:
-        refiner = DiffusionPipeline.from_pretrained(
-            "stabilityai/stable-diffusion-xl-refiner-1.0",
-            vae=vae,
-            torch_dtype=torch.float16,
-            use_safetensors=True,
-            variant="fp16",
-        )
-
-    if ENABLE_CPU_OFFLOAD:
-        pipe.enable_model_cpu_offload()
-        if ENABLE_REFINER:
-            refiner.enable_model_cpu_offload()
-    else:
-        pipe.to(device)
-        if ENABLE_REFINER:
-            refiner.to(device)
-
-    if USE_TORCH_COMPILE:
-        pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
-        if ENABLE_REFINER:
-            refiner.unet = torch.compile(refiner.unet, mode="reduce-overhead", fullgraph=True)
-
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
@@ -80,7 +48,45 @@ def generate(
     num_inference_steps_base: int = 25,
     num_inference_steps_refiner: int = 25,
     apply_refiner: bool = False,
+    model = 'stabilityai/stable-diffusion-xl-base-1.0',
+    vaecall = 'madebyollin/sdxl-vae-fp16-fix',
+    lora = 'pierroromeu/lora-trained-xl-folder',
 ) -> PIL.Image.Image:
+    if torch.cuda.is_available():
+        vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
+        # Substitua o valor dinâmico antes de chamar from_pretrained
+        pipe_model_name = model
+        pipe = DiffusionPipeline.from_pretrained(
+            pipe_model_name,
+            vae=vae,
+            torch_dtype=torch.float16,
+            use_safetensors=True,
+        )
+        if ENABLE_REFINER:
+            refiner_model_name = model
+            refiner = DiffusionPipeline.from_pretrained(
+                refiner_model_name,
+                vae=vae,
+                torch_dtype=torch.float16,
+                use_safetensors=True,
+            )
+
+        pipe.load_lora_weights(lora)
+                                   
+        if ENABLE_CPU_OFFLOAD:
+            pipe.enable_model_cpu_offload()
+            if ENABLE_REFINER:
+                refiner.enable_model_cpu_offload()
+        else:
+            pipe.to(device)
+            if ENABLE_REFINER:
+                refiner.to(device)
+
+        if USE_TORCH_COMPILE:
+            pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+            if ENABLE_REFINER:
+                refiner.unet = torch.compile(refiner.unet, mode="reduce-overhead", fullgraph=True)
+
     generator = torch.Generator().manual_seed(seed)
 
     if not use_negative_prompt:
@@ -142,6 +148,9 @@ with gr.Blocks(css="style.css") as demo:
         visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
     )
     with gr.Group():
+        model = gr.Text(label='Modelo')
+        vaecall = gr.Text(label='VAE')
+        lora = gr.Text(label='LoRA')
         with gr.Row():
             prompt = gr.Text(
                 label="Prompt",
@@ -299,10 +308,13 @@ with gr.Blocks(css="style.css") as demo:
             num_inference_steps_base,
             num_inference_steps_refiner,
             apply_refiner,
+            model,
+            vaecall,
+            lora,
         ],
         outputs=result,
         api_name="run",
     )
 
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
+    demo.queue(max_size=20).launch()