Update raw.py

raw.py

@@ -16,44 +16,41 @@ import gradio as gr
 huggingface_token = os.getenv("HUGGINFACE_TOKEN")
 MAX_SEED = 1000000
 
-quant_config = TransformersBitsAndBytesConfig(load_in_8bit=True,)
-text_encoder_2_8bit = T5EncoderModel.from_pretrained(
-    "LPX55/FLUX.1-merged_uncensored",
-    subfolder="text_encoder_2",
-    quantization_config=quant_config,
-    torch_dtype=torch.bfloat16,
-    token=huggingface_token
-)
+# quant_config = TransformersBitsAndBytesConfig(load_in_8bit=True,)
+# text_encoder_2_8bit = T5EncoderModel.from_pretrained(
+#     "LPX55/FLUX.1-merged_uncensored",
+#     subfolder="text_encoder_2",
+#     quantization_config=quant_config,
+#     torch_dtype=torch.bfloat16,
+#     token=huggingface_token
+# )
 
-good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16, token=huggingface_token).to("cuda")
+# good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16, token=huggingface_token).to("cuda")
 
 # Load pipeline
-controlnet = FluxControlNetModel.from_pretrained(
-    "jasperai/Flux.1-dev-Controlnet-Upscaler",
-    torch_dtype=torch.bfloat16
-)
+# controlnet = FluxControlNetModel.from_pretrained(
+#     "jasperai/Flux.1-dev-Controlnet-Upscaler",
+#     torch_dtype=torch.bfloat16
+# )
 pipe = FluxControlNetPipeline.from_pretrained(
-    "LPX55/FLUX.1-merged_uncensored",
-    controlnet=controlnet,
+    "LPX55/FLUX.1M-8step_upscaler-cnet",
     torch_dtype=torch.bfloat16,
-    vae=good_vae,
-    text_encoder_2=text_encoder_2_8bit,
     token=huggingface_token
 )
-adapter_id = "alimama-creative/FLUX.1-Turbo-Alpha"
-adapter_id2 = "XLabs-AI/flux-RealismLora"
-adapter_id3 = "enhanceaiteam/Flux-uncensored-v2"
+# adapter_id = "alimama-creative/FLUX.1-Turbo-Alpha"
+# adapter_id2 = "XLabs-AI/flux-RealismLora"
+# adapter_id3 = "enhanceaiteam/Flux-uncensored-v2"
 
 pipe.to("cuda")
-pipe.load_lora_weights(adapter_id, adapter_name="turbo")
-pipe.load_lora_weights(adapter_id2, adapter_name="real")
-pipe.load_lora_weights(adapter_id3, weight_name="lora.safetensors", adapter_name="enhance")
-pipe.set_adapters(["turbo", "real", "enhance"], adapter_weights=[0.9, 0.66, 0.6])
-pipe.fuse_lora(adapter_names=["turbo", "real", "enhance"], lora_scale=1.0)
-pipe.unload_lora_weights()
+# pipe.load_lora_weights(adapter_id, adapter_name="turbo")
+# pipe.load_lora_weights(adapter_id2, adapter_name="real")
+# pipe.load_lora_weights(adapter_id3, weight_name="lora.safetensors", adapter_name="enhance")
+# pipe.set_adapters(["turbo", "real", "enhance"], adapter_weights=[0.9, 0.66, 0.6])
+# pipe.fuse_lora(adapter_names=["turbo", "real", "enhance"], lora_scale=1.0)
+# pipe.unload_lora_weights()
 # pipe.enable_xformers_memory_efficient_attention()
 # save to the Hub
-pipe.push_to_hub("FLUX.1M-8step_upscaler-cnet")
+# pipe.push_to_hub("FLUX.1M-8step_upscaler-cnet")
 
 @spaces.GPU
 def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_scale, guidance_scale, seed, guidance_end):