Update app_v5.py

app_v5.py

@@ -1,17 +1,15 @@
-# app_v5.py
+# app_v4.py
 import gradio as gr
-import spaces
-import logging
 import torch
 from gradio_client import Client, handle_file
+import spaces
 import os
 import datetime
 import io
+import numpy as np
 import moondream as md
 from transformers import T5EncoderModel
-from diffusers import FluxControlNetPipeline, FluxPipeline, AutoModel
-from diffusers.quantizers import PipelineQuantizationConfig
-
+from diffusers import FluxControlNetPipeline, FluxControlNetInpaintPipeline, FluxTransformer2DModel
 from diffusers.utils import load_image
 from PIL import Image
 from threading import Thread
@@ -19,12 +17,40 @@ from typing import Generator
 from huggingface_hub import CommitScheduler, HfApi
 from debug import log_params, scheduler, save_image
 from huggingface_hub.utils._runtime import dump_environment_info
+import logging
 
+#############################
+os.environ.setdefault('GRADIO_ANALYTICS_ENABLED', 'False')
+os.environ.setdefault('HF_HUB_DISABLE_TELEMETRY', '1')
 
-logging.basicConfig(level=logging.INFO)
+logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
+#############################
+
+presets = {
+    "Strict Upscale": {
+        "scale": 1.0,
+        "steps": 8,
+        "controlnet_conditioning_scale": 0.75,
+        "guidance_scale": 4.0,
+        "guidance_end": 0.9
+    },
+    "Creative Upscale": {
+        "scale": 2.0,
+        "steps": 6,
+        "controlnet_conditioning_scale": 0.42,
+        "guidance_scale": 3.0,
+        "guidance_end": 0.5
+    },
+    "High Detail Upscale": {
+        "scale": 1.25,
+        "steps": 10,
+        "controlnet_conditioning_scale": 0.9,
+        "guidance_scale": 10.0,
+        "guidance_end": 0.9
+    }
+}
 
-# Ensure device is set
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 MAX_SEED = 1000000
 
@@ -39,66 +65,52 @@ try:
 except Exception as e:
     print(f"Error setting memory usage: {e}")
 
-pipeline_quant_config = PipelineQuantizationConfig(
-    quant_backend="bitsandbytes_4bit",
-    quant_kwargs={"load_in_4bit": True, "bnb_4bit_quant_type": "nf4", "bnb_4bit_compute_dtype": torch.bfloat16},
-    components_to_quantize=["transformer", "text_encoder_2"],
+text_encoder_2_unquant = T5EncoderModel.from_pretrained(
+    "LPX55/FLUX.1-merged_lightning_v2",
+    subfolder="text_encoder_2",
+    torch_dtype=torch.bfloat16,
+    token=huggingface_token
 )
+transformer = FluxTransformer2DModel.from_pretrained(
+        "LPX55/FLUX.1-merged_lightning_v2", subfolder='transformer', torch_dytpe=torch.bfloat16
+    )
+pipe_upscaler = FluxControlNetPipeline.from_pretrained(
+    "LPX55/FLUX.1M-8step_upscaler-cnet",
+    torch_dtype=torch.bfloat16,
+    text_encoder_2=text_encoder_2_unquant,
+    token=huggingface_token
+)
+pipe_upscaler.to("cuda")
 
-# quant_config_5_t5 = TransformersBitsAndBytesConfig(load_in_8bit=True,)
-# text_encoder_2_8b = T5EncoderModel.from_pretrained(
-#     "LPX55/FLUX.1-merged_lightning_v2",
-#     subfolder="text_encoder_2",
-#     quantization_config=quant_config_5_t5,
-#     torch_dtype=torch.float16,
-# )
+controlnet = FluxControlNetModel.from_pretrained("alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta", torch_dtype=torch.bfloat16)
 
-# quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True,)
-# transformer_8bit = FluxPipeline.from_pretrained(
-#     "LPX55/FLUX.1-merged_lightning_v2",
-#     subfolder="transformer",
-#     quantization_config=quant_config,
-#     torch_dtype=torch.float16,
-# )
-pipe = FluxControlNetPipeline.from_pretrained(
-    "LPX55/FLUX.1M-8step_upscaler-cnet",
-    quantization_config=pipeline_quant_config,
-    torch_dtype=torch.float16
+pipe = FluxControlNetInpaintPipeline.from_pretrained(
+    "LPX55/FLUX.1-merged_lightning_v2",
+    controlnet=controlnet,
+    transformer=transformer,
+    torch_dtype=torch.bfloat16,
+    token=huggingface_token
 )
 pipe.to("cuda")
+pipe.transformer.to(torch.bfloat16)
+pipe.controlnet.to(torch.bfloat16)
 
 try:
     dump_environment_info()
 except Exception as e:
     print(f"Failed to dump env info: {e}")
+    
+def get_image_dimensions(image: Image) -> str:
+    width, height = image.size
+    return f"Original Image Dimensions: {width}x{height}"
+    
+def resize_image_to_max_side(image: Image, max_side_length=1024) -> Image:
+    width, height = image.size
+    ratio = min(max_side_length / width, max_side_length / height)
+    new_size = (int(width * ratio), int(height * ratio))
+    resized_image = image.resize(new_size, Image.BILINEAR)
+    return resized_image
 
-@spaces.GPU(duration=6, progress=gr.Progress(track_tqdm=True))
-@torch.no_grad()
-def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_scale, guidance_scale, seed, guidance_end):
-    generator = torch.Generator().manual_seed(seed)
-    # Load control image
-    control_image = load_image(control_image)
-    w, h = control_image.size
-    w = w - w % 32
-    h = h - h % 32
-    control_image = control_image.resize((int(w * scale), int(h * scale)), resample=2)  # Resample.BILINEAR
-    print("Size to: " + str(control_image.size[0]) + ", " + str(control_image.size[1]))
-    print(f"PromptLog: {repr(prompt)}")
-    with torch.inference_mode():
-        image = pipe(
-            generator=generator,
-            prompt=prompt,
-            control_image=control_image,
-            controlnet_conditioning_scale=controlnet_conditioning_scale,
-            num_inference_steps=steps,
-            guidance_scale=guidance_scale,
-            height=control_image.size[1],
-            width=control_image.size[0],
-            control_guidance_start=0.0,
-            control_guidance_end=guidance_end,
-        ).images[0]
-        # print("Type: " + str(type(image)))
-    return image
 
 def combine_caption_focus(caption, focus):
     try:
@@ -110,55 +122,138 @@ def combine_caption_focus(caption, focus):
     except Exception as e:
         print(f"Error combining caption and focus: {e}")
         return "highly detailed photo, raw photography."
-
 def generate_caption(control_image):
     try:
         if control_image is None:
-            return "Waiting for control image..."
+            return "Waiting for control image...", "Original Image Dimensions: N/A"
+           
+        # Get original dimensions
+        original_dimensions = get_image_dimensions(control_image)
         
+        # Resize the image to a maximum longest side of 1024 pixels
+        control_image = resize_image_to_max_side(control_image, max_side_length=1024)
+
         # Generate a detailed caption
         mcaption = model.caption(control_image, length="short")
         detailed_caption = mcaption["caption"]
         print(f"Detailed caption: {detailed_caption}")
         
-        return detailed_caption
+        return detailed_caption, original_dimensions
     except Exception as e:
         print(f"Error generating caption: {e}")
-        return "A detailed photograph"
+        return "A detailed photograph", "Original Image Dimensions: N/A"
 
 def generate_focus(control_image, focus_list):
     try:
         if control_image is None:
-            return None
+            return None, "Original Image Dimensions: N/A"
         if focus_list is None:
-            return ""
+            return "", "Original Image Dimensions: N/A"
+        
+        # Get original dimensions
+        original_dimensions = get_image_dimensions(control_image)
+        
+        # Resize the image to a maximum longest side of 1024 pixels
+        control_image = resize_image_to_max_side(control_image, max_side_length=1024)
+
         # Generate a detailed caption
         focus_query = model.query(control_image, "Please provide a concise but illustrative description of the following area(s) of focus: " + focus_list)
         focus_description = focus_query["answer"]
         print(f"Areas of focus: {focus_description}") 
-
-        return focus_description
+        return focus_description, original_dimensions
     except Exception as e:
         print(f"Error generating focus: {e}")
-        return "highly detailed photo, raw photography."
+        return "highly detailed photo, raw photography.", "Original Image Dimensions: N/A"
+        
+
+@spaces.GPU(duration=6, progress=gr.Progress(track_tqdm=True))
+@torch.no_grad()
+def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_scale, guidance_scale, seed, guidance_end):
+    generator = torch.Generator().manual_seed(seed)
+    
+    # Ensure transparency is preserved
+    control_image = control_image.convert("RGBA")
+    
+    # Resize the image to a maximum longest side of 1024 pixels
+    control_image = resize_image_to_max_side(control_image, max_side_length=1024)
+    
+    w, h = control_image.size
+    # Crop to nearest multiple of 32
+    w = w - w % 32
+    h = h - h % 32
+    # Corrected resizing code
+    control_image = control_image.resize((w, h), resample=2)
+    
+    print(f"Resized image dimensions: {control_image.size[0]}x{control_image.size[1]}")
+    print(f"PromptLog: {repr(prompt)}")
+    
+    # Convert image to RGB for processing, but keep alpha channel for transparency
+    control_image_rgb = control_image.convert("RGB")
+    control_image_alpha = control_image.split()[-1]
+    
+    # Convert alpha channel to a mask (transparent = white, opaque = black)
+    # White corresponds to 1 (to be inpainted), black corresponds to 0 (to be preserved)
+    # Convert alpha to numpy array for processing
+    alpha_array = np.array(control_image_alpha)
+    # Create binary mask (1 for transparent, 0 for opaque)
+    mask = (alpha_array > 128).astype(np.float32)  # 1 for transparent (to be inpainted), 0 for opaque
+    
+    # Optional: Visualize the mask (for debugging purposes)
+    # mask_image = Image.fromarray((mask * 255).astype(np.uint8))
+    # mask_image.show()
+    
+    with torch.inference_mode():
+        image = pipe(
+            image=control_image_rgb,
+            generator=generator,
+            prompt=prompt,
+            control_image=control_image_rgb,
+            mask_image=mask,  # Pass the numpy array as the mask
+            controlnet_conditioning_scale=controlnet_conditioning_scale,
+            num_inference_steps=steps,
+            guidance_scale=guidance_scale,
+            height=control_image.size[1],
+            width=control_image.size[0],
+            control_guidance_start=0.0,
+            control_guidance_end=guidance_end,
+        ).images[0]
+        
+        # Reapply the alpha channel to the generated image
+        image = image.convert("RGBA")
+        image.putalpha(control_image_alpha)
+        
+        return image
+    
+def update_parameters(preset):
+    if preset in presets:
+        params = presets[preset]
+        return (
+            params["scale"],
+            params["steps"],
+            params["controlnet_conditioning_scale"],
+            params["guidance_scale"],
+            params["guidance_end"]
+        )
+    else:
+        # Default values if preset is not found
+        return 1.0, 8, 0.6, 3.5, 1.0
 
 def process_image(control_image, user_prompt, system_prompt, scale, steps, 
                 controlnet_conditioning_scale, guidance_scale, seed, 
-                guidance_end, temperature, top_p, max_new_tokens, log_prompt):
+                guidance_end, temperature, max_new_tokens, log_prompt):
     # Initialize with empty caption
     final_prompt = user_prompt.strip()
     # If no user prompt provided, generate a caption first
     if not final_prompt:
         # Generate a detailed caption
-        print("Generating caption...")
+
         mcaption = model.caption(control_image, length="normal")
         detailed_caption = mcaption["caption"]
         final_prompt = detailed_caption
         yield f"Using caption: {final_prompt}", None, final_prompt
-    
-    # Show the final prompt being used
+
     yield f"Generating with: {final_prompt}", None, final_prompt
-    
+
     # Generate the image
     try:
         image = generate_image(
@@ -171,10 +266,12 @@ def process_image(control_image, user_prompt, system_prompt, scale, steps,
             seed=seed,
             guidance_end=guidance_end
         )
-        
         try:
-            debug_img = Image.open(image.save("/tmp/" + str(seed) + "output.png"))
-            save_image("/tmp/" + str(seed) + "output.png", debug_img)
+            # Ensure the image is saved with transparency
+            with io.BytesIO() as output:
+                image.save(output, format="PNG")
+                debug_img = Image.open(output).convert("RGBA")
+                save_image("/tmp/" + str(seed) + "output.png", debug_img)
         except Exception as e:
             print("Error 160: " + str(e))
         log_params(final_prompt, scale, steps, controlnet_conditioning_scale, guidance_scale, seed, guidance_end, control_image, image)
@@ -182,7 +279,7 @@ def process_image(control_image, user_prompt, system_prompt, scale, steps,
     except Exception as e:
         print("Error: " + str(e))
         yield f"Error: {str(e)}", None, None
-    
+
 with gr.Blocks(title="FLUX Turbo Upscaler", fill_height=True) as demo:
     gr.Markdown("⚠️ WIP SPACE - UNFINISHED & BUGGY")
     # status_box = gr.Markdown("🔄 Warming up...")
@@ -196,16 +293,21 @@ with gr.Blocks(title="FLUX Turbo Upscaler", fill_height=True) as demo:
         with gr.Column(scale=1):
             prompt = gr.Textbox(lines=4, info="Enter your prompt here or wait for auto-generation...", label="Image Description")
             focus = gr.Textbox(label="Area(s) of Focus", info="e.g. 'face', 'eyes', 'hair', 'clothes', 'background', etc.", value="clothing material, textures, ethnicity")
-            scale = gr.Slider(1, 3, value=1, label="Scale (Upscale Factor)", step=0.25)
+            scale = gr.Slider(1, 3, value=1, label="Scale (Upscale Factor)", step=0.1)
             with gr.Row():
                 generate_button = gr.Button("Generate Image", variant="primary")
                 caption_button = gr.Button("Generate Caption", variant="secondary")
         with gr.Column(scale=1):
+            with gr.Row():
+                preset_choices = list(presets.keys())
+                preset_radio = gr.Radio(choices=preset_choices, label="Select Preset", value=preset_choices[0])
             seed = gr.Slider(0, MAX_SEED, value=42, label="Seed", step=1)
             steps = gr.Slider(2, 16, value=8, label="Steps", step=1)
             controlnet_conditioning_scale = gr.Slider(0, 1, value=0.6, label="ControlNet 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")
+            original_dimensions = gr.Markdown(value="Original Image Dimensions: N/A")  # New output for dimensions
+
     with gr.Row():
         with gr.Accordion("Auto-Caption settings", open=False, visible=False):
             system_prompt = gr.Textbox(
@@ -220,11 +322,6 @@ with gr.Blocks(title="FLUX Turbo Upscaler", fill_height=True) as demo:
                 info="Higher values make the output more random, lower values make it more deterministic.",
                 visible=False  # Changed to visible
             )
-            top_p_slider = gr.Slider(
-                minimum=0.0, maximum=1.0, value=0.9, step=0.01,
-                label="Top-p",
-                visible=False  # Changed to visible
-            )
             max_tokens_slider = gr.Slider(
                 minimum=1, maximum=2048, value=368, step=1,
                 label="Max New Tokens",
@@ -234,65 +331,39 @@ with gr.Blocks(title="FLUX Turbo Upscaler", fill_height=True) as demo:
         log_prompt = gr.Checkbox(value=True, label="Log", visible=False)  # Changed to visible
     
     gr.Markdown("**Tips:** 8 steps is all you need! Incredibly powerful tool, usage instructions coming soon.")
-    with gr.Accordion("Help,I keep getting ZeroGPU errors.", open=False, elem_id="zgpu"):
+    with gr.Accordion("Auth for those Getting ZeroGPU errors.", open=False, elem_id="zgpu"):
         msg1 = gr.Markdown()
         try_btn = gr.LoginButton()
-        try:
-            x_ip_token = request.headers['x-ip-token']
-            client = Client("LPX55/zerogpu-experiments", hf_token=huggingface_token, headers={"x-ip-token": x_ip_token})
-            cresult = client.predict(
-            		n=3,
-            		api_name="/predict"
-            )
-            print(f"X TOKEN: {x_ip_token}")
-            print(cresult)
-        except:
-            print("Guess we're just going to have to pretend that Spaces have been broken for almost a year now..")
         
-        # result = client.predict(
-        # 		image=handle_file('https://raw.githubusercontent.com/gradio-app/gradio/main/test/test_files/bus.png'),
-        # 		width=1024,
-        # 		height=1024,
-        # 		overlap_percentage=10,
-        # 		num_inference_steps=8,
-        # 		resize_option="Full",
-        # 		custom_resize_percentage=50,
-        # 		prompt_input="Hello!!",
-        # 		alignment="Middle",
-        # 		overlap_left=True,
-        # 		overlap_right=True,
-        # 		overlap_top=True,
-        # 		overlap_bottom=True,
-        # 		x_offset=0,
-        # 		y_offset=0,
-        # 		api_name="/infer"
+        # sus = ['x-zerogpu-token', 'x-zerogpu-uuid', 'x-proxied-host', 'x-proxied-path', 'x-proxied-replica', 'x-request-id', 'x-ip-token']
+        # x_ip_token = request.headers['X-IP-TOKEN']
+        # print(str(x_ip_token))
+        # client = Client("LPX55/zerogpu-experiments", hf_token=huggingface_token, headers={"x-ip-token": x_ip_token})
+        # cresult = client.predict(
+        #         n=3,
+        #         api_name="/predict"
         # )
-    def hello(profile: gr.OAuthProfile | None) -> str:
-        if profile is None:
-            return "Hello guest! There is a bug with HF ZeroGPUs that are afffecting some usage on certain spaces. Testing out some possible solutions."
-        return f"You are logged in as {profile.name}. If you run into incorrect messages about ZeroGPU runtime credits being out, PLEASE give me a heads up so I can investigate further."
 
     caption_state = gr.State()
     focus_state = gr.State()
     log_state = gr.State()
-
     generate_button.click(
         fn=process_image,
         inputs=[
             control_image, prompt, system_prompt, scale, steps, 
             controlnet_conditioning_scale, guidance_scale, seed, 
-            guidance_end, temperature_slider, top_p_slider, max_tokens_slider, log_prompt
+            guidance_end, temperature_slider, max_tokens_slider, log_prompt
         ],
         outputs=[log_state, generated_image, prompt]
     )
     control_image.upload(
         generate_caption,
         inputs=[control_image],
-        outputs=[caption_state]
+        outputs=[caption_state, original_dimensions]
     ).then(
         generate_focus,
         inputs=[control_image, focus],
-        outputs=[focus_state]
+        outputs=[focus_state, original_dimensions]
     ).then(
         combine_caption_focus,
         inputs=[caption_state, focus_state],
@@ -301,15 +372,25 @@ with gr.Blocks(title="FLUX Turbo Upscaler", fill_height=True) as demo:
     caption_button.click(
         fn=generate_caption,
         inputs=[control_image],
-        outputs=[prompt]
+        outputs=[prompt, original_dimensions]
     ).then(
         generate_focus,
         inputs=[control_image, focus],
-        outputs=[focus_state]
+        outputs=[focus_state, original_dimensions]
     ).then(
         combine_caption_focus,
         inputs=[caption_state, focus_state],
         outputs=[prompt]
     )
+    preset_radio.change(
+        fn=update_parameters,
+        inputs=[preset_radio],
+        outputs=[scale, steps, controlnet_conditioning_scale, guidance_scale, guidance_end]
+    )
+    def hello(profile: gr.OAuthProfile | None) -> str:
+        if profile is None:
+            return "Hello guest! There is a bug with HF ZeroGPUs that are afffecting some usage on certain spaces. Testing out some possible solutions."
+        return f"You are logged in as {profile.name}. If you run into incorrect messages about ZeroGPU runtime credits being out, PLEASE give me a heads up so I can investigate further."
+
     demo.load(hello, inputs=None, outputs=msg1)
-demo.queue().launch(show_error=True)
+demo.queue().launch(show_error=True)