Update app.py

app.py

@@ -7,7 +7,7 @@ import torch
 import spaces
 import gradio as gr
 
-from diffusers import FluxInpaintPipeline
+from diffusers import FluxTransformer2DModel, FluxInpaintPipeline
 import google.generativeai as genai
 
 
@@ -49,14 +49,19 @@ def gemini_predict(prompt):
 MAX_SEED = np.iinfo(np.int32).max
 DEVICE = "cuda" #if torch.cuda.is_available() else "cpu"
 
-#Setting up Flux (Schnell) Inpainting, with Realism LoRA weights
-inpaint_pipe = FluxInpaintPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to(DEVICE)
-inpaint_pipe.load_lora_weights('hugovntr/flux-schnell-realism', weight_name='schnell-realism_v2.3.safetensors', adapter_name="better")
-inpaint_pipe.set_adapters(["better"], adapter_weights=[2.6])
-inpaint_pipe.fuse_lora(adapter_name=["better"], lora_scale=1.0)
-inpaint_pipe.unload_lora_weights()
+#Setting up Flux (Schnell) Inpainting
+#inpaint_pipe = FluxInpaintPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to(DEVICE)
 
-torch.cuda.empty_cache()
+transformer_SchnellReal = FluxTransformer2DModel.from_single_file("https://huggingface.co/SG161222/RealFlux_1.0b_Schnell/blob/main/4%20-%20Schnell%20Transformer%20Version/RealFlux_1.0b_Schnell_Transformer.safetensors", torch_dtype=torch.bfloat16)
+inpaint_pipe = FluxInpaintPipeline.from_pretrained(bfl_repo, transformer=transformer_SchnellReal, torch_dtype=dtype).to(DEVICE)
+
+#Uncomment the following 4 lines, if you want LoRA Realism weights added to the pipeline
+# inpaint_pipe.load_lora_weights('hugovntr/flux-schnell-realism', weight_name='schnell-realism_v2.3.safetensors', adapter_name="better")
+# inpaint_pipe.set_adapters(["better"], adapter_weights=[2.6])
+# inpaint_pipe.fuse_lora(adapter_name=["better"], lora_scale=1.0)
+# inpaint_pipe.unload_lora_weights()
+
+#torch.cuda.empty_cache()
 
 @spaces.GPU()
 def process(input_image_editor, mask_image, input_text, strength, seed, randomize_seed, num_inference_steps, guidance_scale=3.5, progress=gr.Progress(track_tqdm=True)):