Update app.py

app.py

@@ -9,7 +9,7 @@ import uuid
 import json
 import os
 
-from diffusers import StableDiffusionXLPipeline, StableDiffusion3Pipeline, SD3Transformer2DModel
+from diffusers import StableDiffusionXLPipeline, StableDiffusion3Pipeline, SD3Transformer2DModel, FlashFlowMatchEulerDiscreteScheduler
 from huggingface_hub import snapshot_download
 from peft import PeftModel
 from PIL import Image
@@ -32,10 +32,6 @@ DEFAULT_WIDTH = 1024
 
 headers = {"Content-Type": "application/json", "Authorization": f"Bearer {HF_TOKEN}" }
 
-repo_large_path = snapshot_download(repo_id="stabilityai/stable-diffusion-3-medium", revision="refs/pr/26", token=HF_TOKEN)
-repo_large_transformer_path = SD3Transformer2DModel.from_pretrained(repo_large_path, subfolder="transformer", torch_dtype=torch.float16)
-repo_large_transformer = PeftModel.from_pretrained(repo_large_transformer_path, "jasperai/flash-sd3")
-
 css = '''
 .gradio-container{max-width: 560px !important}
 h1{text-align:center}
@@ -44,21 +40,36 @@ footer {
 }
 '''
 
+'''
 repo_default = StableDiffusionXLPipeline.from_pretrained("fluently/Fluently-XL-Final", torch_dtype=torch.float16, use_safetensors=True, add_watermarker=False)
 repo_default.load_lora_weights("ehristoforu/dalle-3-xl-v2", adapter_name="base")
 repo_default.set_adapters(["base"], adapter_weights=[0.7])
+'''
 
+'''
 repo_pixel = StableDiffusionXLPipeline.from_pretrained("fluently/Fluently-XL-Final", torch_dtype=torch.float16, use_safetensors=True, add_watermarker=False)
 repo_pixel.load_lora_weights("artificialguybr/PixelArtRedmond", adapter_name="base")
 repo_pixel.load_lora_weights("nerijs/pixel-art-xl", adapter_name="base2")
 repo_pixel.set_adapters(["base", "base2"], adapter_weights=[1, 1])
+'''
+
+repo_large_path = snapshot_download(repo_id="stabilityai/stable-diffusion-3-medium", revision="refs/pr/26", token=HF_TOKEN)
+repo_large_transformer_path = SD3Transformer2DModel.from_pretrained(repo_large_path, subfolder="transformer", torch_dtype=torch.float16)
+repo_large_transformer = PeftModel.from_pretrained(repo_large_transformer_path, "jasperai/flash-sd3")
+
+pipe.scheduler = FlashFlowMatchEulerDiscreteScheduler.from_pretrained(
+  repo_large_path,
+  subfolder="scheduler",
+)
 
 repo_customs = {
-    "Default": repo_default,
-    "Realistic": StableDiffusionXLPipeline.from_pretrained("ehristoforu/Visionix-alpha", torch_dtype=torch.float16, use_safetensors=True, add_watermarker=False),
-    "Anime": StableDiffusionXLPipeline.from_pretrained("cagliostrolab/animagine-xl-3.1", torch_dtype=torch.float16, use_safetensors=True, add_watermarker=False),
-    "Pixel": repo_pixel,
-    "Large": StableDiffusion3Pipeline.from_pretrained(repo_large_path, transformer=repo_large_transformer, torch_dtype=torch.float16, use_safetensors=True),
+    "Default": None, #repo_default,
+    "Realistic": None, #StableDiffusionXLPipeline.from_pretrained("ehristoforu/Visionix-alpha", torch_dtype=torch.float16, use_safetensors=True, add_watermarker=False),
+    "Anime": None, #StableDiffusionXLPipeline.from_pretrained("cagliostrolab/animagine-xl-3.1", torch_dtype=torch.float16, use_safetensors=True, add_watermarker=False),
+    "Pixel": None, #repo_pixel,
+    "Large": StableDiffusion3Pipeline.from_pretrained(repo_large_path, transformer=repo_large_transformer, torch_dtype=torch.float16, use_safetensors=True,
+        text_encoder_3=None,
+        tokenizer_3=None),
 }
 
 # Functions
@@ -120,13 +131,13 @@ def generate(input=DEFAULT_INPUT, filter_input="", negative_input=DEFAULT_NEGATI
     parameters  = {
         "prompt": input,
         "negative_prompt": filter_input + negative_input,
-        "height": height,
-        "width": width,
+        #"height": height,
+        #"width": width,
         "num_inference_steps": steps,
         "guidance_scale": guidance,
-        "num_images_per_prompt": number,
+        #"num_images_per_prompt": number,
         "generator": torch.Generator().manual_seed(seed),
-        "output_type":"pil",
+        #"output_type":"pil",
     }
     
     images = repo(**parameters).images