Update app.py

app.py

@@ -17,6 +17,7 @@ from diffusers import DiffusionPipeline, FluxTransformer2DModel, FluxPipeline, A
 import safetensors.torch
 from safetensors.torch import load_file
 from custom_pipeline import FluxWithCFGPipeline
+from transformers import CLIPModel, CLIPProcessor, CLIPConfig
 import gc
 
 cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
@@ -29,12 +30,27 @@ torch.backends.cuda.matmul.allow_tf32 = True
 dtype = torch.float16
 pipe = FluxWithCFGPipeline.from_pretrained(
     "ostris/OpenFLUX.1", torch_dtype=dtype
-)
-# pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype)
-pipe.load_lora_weights("ostris/OpenFLUX.1", weight_name="openflux1-v0.1.0-fast-lora.safetensors", adapter_name="fast")
-pipe.set_adapters("fast")
-pipe.fuse_lora(adapter_names=["fast"], lora_scale=1.0)
+).to("cuda")
+pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to("cuda")
+
 pipe.to("cuda")
+clipmodel = 'norm'
+if clipmodel == "long":
+    model_id = "zer0int/LongCLIP-GmP-ViT-L-14"
+    config = CLIPConfig.from_pretrained(model_id)
+    maxtokens = 77
+if clipmodel == "norm":
+    model_id = "zer0int/CLIP-GmP-ViT-L-14"
+    config = CLIPConfig.from_pretrained(model_id)
+    maxtokens = 77
+clip_model = CLIPModel.from_pretrained(model_id, torch_dtype=torch.bfloat16, config=config, ignore_mismatched_sizes=True).to("cuda")
+clip_processor = CLIPProcessor.from_pretrained(model_id, padding="max_length", max_length=maxtokens, ignore_mismatched_sizes=True, return_tensors="pt", truncation=True)
+config.text_config.max_position_embeddings = 77
+
+pipe.tokenizer = clip_processor.tokenizer
+pipe.text_encoder = clip_model.text_model
+pipe.tokenizer_max_length = maxtokens
+pipe.text_encoder.dtype = torch.bfloat16
 torch.cuda.empty_cache()
 
 # Load LoRAs from JSON file
@@ -119,9 +135,13 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, wid
     # Load LoRA weights
     with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"):
         if "weights" in selected_lora:
-            pipe.load_lora_weights(lora_path, weight_name=selected_lora["weights"])
+            pipe.load_lora_weights(lora_path, weight_name=selected_lora["weights"], adapter_name="choice")
+            pipe.load_lora_weights("ostris/OpenFLUX.1", weight_name="openflux1-v0.1.0-fast-lora.safetensors", adapter_name="fast")
+            pipe.set_adapters(["fast", "choice"], adapter_weights=[1.0, lora_scale])
         else:
-            pipe.load_lora_weights(lora_path)
+            pipe.load_lora_weights(lora_path, adapter_name="choice")
+            pipe.load_lora_weights("ostris/OpenFLUX.1", weight_name="openflux1-v0.1.0-fast-lora.safetensors", adapter_name="fast")
+            pipe.set_adapters(["fast", "choice"], adapter_weights=[1.0, lora_scale])
         
     # Set random seed for reproducibility
     with calculateDuration("Randomizing seed"):
@@ -134,6 +154,10 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, wid
     return image, seed  
 
 run_lora.zerogpu = True
+#pipe.load_lora_weights("ostris/OpenFLUX.1", weight_name="openflux1-v0.1.0-fast-lora.safetensors", adapter_name="fast")
+#pipe.set_adapters("fast")
+#pipe.set_adapters(["fast", "toy"], adapter_weights=[0.5, 1.0])
+#pipe.fuse_lora(adapter_names=["fast"], lora_scale=1.0)
 
 css = '''
 #gen_btn{height: 100%}
@@ -180,8 +204,8 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as app:
         with gr.Accordion("Advanced Settings", open=True):
             with gr.Column():
                 with gr.Row():
-                    cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=1.0)
-                    steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=5)
+                    cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.0)
+                    steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=10)
                 
                 with gr.Row():
                     width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=768)