Update raw.py

raw.py

@@ -2,6 +2,7 @@ import torch
 import spaces
 import os
 from diffusers.utils import load_image
+from diffusers.hooks import apply_group_offloading
 from diffusers import FluxControlNetModel, FluxControlNetPipeline, AutoencoderKL
 from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
 from transformers import T5EncoderModel
@@ -43,13 +44,62 @@ pipe = FluxControlNetPipeline.from_pretrained(
 # adapter_id3 = "enhanceaiteam/Flux-uncensored-v2"
 
 pipe.to("cuda")
+try:
+    apply_group_offloading(
+        pipe.transformer,
+        offload_type="leaf_level",
+        offload_device=torch.device("cpu"),
+        onload_device=torch.device("cuda"),
+        use_stream=True,
+    )
+    apply_group_offloading(
+        pipe.text_encoder, 
+        offload_device=torch.device("cpu"),
+        onload_device=torch.device("cuda"),
+        offload_type="leaf_level",
+        use_stream=True,
+    )
+    apply_group_offloading(
+        pipe.text_encoder_2, 
+        offload_device=torch.device("cpu"),
+        onload_device=torch.device("cuda"),
+        offload_type="leaf_level",
+        use_stream=True,
+    )
+    apply_group_offloading(
+        pipe.vae, 
+        offload_device=torch.device("cpu"),
+        onload_device=torch.device("cuda"),
+        offload_type="leaf_level",
+        use_stream=True,
+    )
+except:
+    console.log("debug-group")
+
+try:
+    pipe.enable_sequential_cpu_offload()
+except:
+    console.log("debug-1")
+try:
+    pipe.vae.enable_slicing()
+except:
+    console.log("debug-2")
+try:
+    pipe.vae.enable_tiling()
+except:
+    console.log("debug-3")
+try:
+    pipe.enable_xformers_memory_efficient_attention()
+except:
+    console.log("debug-4")
+
+
 # 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")
 
@@ -74,7 +124,9 @@ def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_
             guidance_scale=guidance_scale,
             height=control_image.size[1],
             width=control_image.size[0],
-            control_guidance_end=guidance_end
+            control_guidance_start=0.0,
+            control_guidance_end=guidance_end,
+            guidance_scale=30.0,
         ).images[0]
         
     return image
@@ -94,7 +146,7 @@ with gr.Blocks(title="FLUX Turbo Upscaler", fill_height=True) as iface:
             seed = gr.Slider(0, MAX_SEED, value=42, label="Seed", step=1)
             steps = gr.Slider(2, 16, value=8, label="Steps")
             controlnet_conditioning_scale = gr.Slider(0, 1, value=0.6, label="ControlNet Scale")
-            guidance_scale = gr.Slider(1, 20, value=3.5, label="Guidance Scale")
+            guidance_scale = gr.Slider(1, 30, value=3.5, label="Guidance Scale")
             guidance_end = gr.Slider(0, 1, value=1.0, label="Guidance End")