Update app.py

app.py

@@ -6,69 +6,20 @@ import torch
 import spaces
 import random
 
-from diffusers import FluxFillPipeline
+from diffusers import FluxPipeline
 from PIL import Image
 
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
-pipe = FluxFillPipeline.from_pretrained("Himanshu806/FluxHyperReal", torch_dtype=torch.bfloat16).to("cuda")
-# pipe.enable_sequential_cpu_offload()
-# pipe.enable_fp16()
-# pipe.vae.enable_slicing()
-# pipe.vae.enable_tiling()
-
-def calculate_optimal_dimensions(image: Image.Image):
-    # Extract the original dimensions
-    original_width, original_height = image.size
-
-    # Set constants
-    MIN_ASPECT_RATIO = 9 / 16
-    MAX_ASPECT_RATIO = 16 / 9
-    FIXED_DIMENSION = 1024
-
-    # Calculate the aspect ratio of the original image
-    original_aspect_ratio = original_width / original_height
-
-    # Determine which dimension to fix
-    if original_aspect_ratio > 1:  # Wider than tall
-        width = FIXED_DIMENSION
-        height = round(FIXED_DIMENSION / original_aspect_ratio)
-    else:  # Taller than wide
-        height = FIXED_DIMENSION
-        width = round(FIXED_DIMENSION * original_aspect_ratio)
-
-    # Ensure dimensions are multiples of 8
-    width = (width // 8) * 8
-    height = (height // 8) * 8
-
-    # Enforce aspect ratio limits
-    calculated_aspect_ratio = width / height
-    if calculated_aspect_ratio > MAX_ASPECT_RATIO:
-        width = (height * MAX_ASPECT_RATIO // 8) * 8
-    elif calculated_aspect_ratio < MIN_ASPECT_RATIO:
-        height = (width / MIN_ASPECT_RATIO // 8) * 8
-
-    # Ensure width and height remain above the minimum dimensions
-    width = max(width, 576) if width == FIXED_DIMENSION else width
-    height = max(height, 576) if height == FIXED_DIMENSION else height
-
-    return width, height
+pipe = FluxPipeline.from_pretrained("Himanshu806/FluxHyperReal", torch_dtype=torch.bfloat16).to("cuda")
 
 @spaces.GPU(durations=300)
-def infer(edit_images, prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
-    # pipe.enable_xformers_memory_efficient_attention()
-
-    image = edit_images["background"]
-    width, height = calculate_optimal_dimensions(image)
-    mask = edit_images["layers"][0]
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
+def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
+    
     image = pipe(
         prompt=prompt,
-        image=image,
-        mask_image=mask,
         height=height,
         width=width,
         guidance_scale=guidance_scale,
@@ -104,15 +55,6 @@ with gr.Blocks(css=css) as demo:
         """)
         with gr.Row():
             with gr.Column():
-                edit_image = gr.ImageEditor(
-                    label='Upload and draw mask for inpainting',
-                    type='pil',
-                    sources=["upload", "webcam"],
-                    image_mode='RGB',
-                    layers=False,
-                    brush=gr.Brush(colors=["#FFFFFF"]),
-                    # height=600
-                )
                 prompt = gr.Text(
                     label="Prompt",
                     show_label=False,
@@ -177,270 +119,8 @@ with gr.Blocks(css=css) as demo:
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn = infer,
-        inputs = [edit_image, prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
+        inputs = [prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
         outputs = [result, seed]
     )
 
-demo.launch()
-
-
-# import gradio as gr
-# import numpy as np
-# import torch
-# import random
-# from PIL import Image
-# import cv2
-# import spaces
-
-# # ------------------ Inpainting Pipeline Setup ------------------ #
-# from diffusers import FluxFillPipeline
-
-# MAX_SEED = np.iinfo(np.int32).max
-# MAX_IMAGE_SIZE = 2048
-
-# pipe = FluxFillPipeline.from_pretrained(
-#     "black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16
-# )
-# pipe.load_lora_weights("alvdansen/flux-koda")
-# pipe.enable_lora()
-
-# def calculate_optimal_dimensions(image: Image.Image):
-#     # Extract the original dimensions
-#     original_width, original_height = image.size
-
-#     # Set constants
-#     MIN_ASPECT_RATIO = 9 / 16
-#     MAX_ASPECT_RATIO = 16 / 9
-#     FIXED_DIMENSION = 1024
-
-#     # Calculate the aspect ratio of the original image
-#     original_aspect_ratio = original_width / original_height
-
-#     # Determine which dimension to fix
-#     if original_aspect_ratio > 1:  # Wider than tall
-#         width = FIXED_DIMENSION
-#         height = round(FIXED_DIMENSION / original_aspect_ratio)
-#     else:  # Taller than wide
-#         height = FIXED_DIMENSION
-#         width = round(FIXED_DIMENSION * original_aspect_ratio)
-
-#     # Ensure dimensions are multiples of 8
-#     width = (width // 8) * 8
-#     height = (height // 8) * 8
-
-#     # Enforce aspect ratio limits
-#     calculated_aspect_ratio = width / height
-#     if calculated_aspect_ratio > MAX_ASPECT_RATIO:
-#         width = (height * MAX_ASPECT_RATIO // 8) * 8
-#     elif calculated_aspect_ratio < MIN_ASPECT_RATIO:
-#         height = (width / MIN_ASPECT_RATIO // 8) * 8
-
-#     # Ensure minimum dimensions are met
-#     width = max(width, 576) if width == FIXED_DIMENSION else width
-#     height = max(height, 576) if height == FIXED_DIMENSION else height
-
-#     return width, height
-
-# # ------------------ SAM (Transformers) Imports and Initialization ------------------ #
-# from transformers import SamModel, SamProcessor
-
-# # Load the model and processor from Hugging Face.
-# sam_model = SamModel.from_pretrained("facebook/sam-vit-base")
-# sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
-
-# @spaces.GPU(durations=300)
-# def generate_mask_with_sam(image: Image.Image, mask_prompt: str):
-#     """
-#     Generate a segmentation mask using SAM (via Hugging Face Transformers).
-
-#     The mask_prompt is expected to be a comma-separated string of two integers,
-#     e.g. "450,600" representing an (x,y) coordinate in the image.
-    
-#     The function converts the coordinate into the proper input format for SAM and returns a binary mask.
-#     """
-#     if mask_prompt.strip() == "":
-#         raise ValueError("No mask prompt provided.")
-    
-#     try:
-#         # Parse the mask_prompt into a coordinate
-#         coords = [int(x.strip()) for x in mask_prompt.split(",")]
-#         if len(coords) != 2:
-#             raise ValueError("Expected two comma-separated integers (x,y).")
-#     except Exception as e:
-#         raise ValueError("Invalid mask prompt. Please provide coordinates as 'x,y'. Error: " + str(e))
-    
-#     # The SAM processor expects a list of input points.
-#     # Format the point as a list of lists; here we assume one point per image.
-#     # (The Transformers SAM expects the points in [x, y] order.)
-#     input_points = [coords]  # e.g. [[450,600]]
-#     # Optionally, you can supply input_labels (1 for foreground, 0 for background)
-#     input_labels = [1]
-    
-#     # Prepare the inputs for the SAM processor.
-#     inputs = sam_processor(images=image,
-#                            input_points=[input_points],
-#                            input_labels=[input_labels],
-#                            return_tensors="pt")
-    
-#     # Move tensors to the same device as the model.
-#     device = next(sam_model.parameters()).device
-#     inputs = {k: v.to(device) for k, v in inputs.items()}
-    
-#     # Forward pass through SAM.
-#     with torch.no_grad():
-#         outputs = sam_model(**inputs)
-    
-#     # The output contains predicted masks; we take the first mask from the first prompt.
-#     # (Assuming outputs.pred_masks is of shape (batch_size, num_masks, H, W))
-#     pred_masks = outputs.pred_masks  # Tensor of shape (1, num_masks, H, W)
-#     mask = pred_masks[0][0].detach().cpu().numpy()
-    
-#     # Convert the mask to binary (0 or 255) using a threshold.
-#     mask_bin = (mask > 0.5).astype(np.uint8) * 255
-#     mask_pil = Image.fromarray(mask_bin)
-#     return mask_pil
-
-# # ------------------ Inference Function ------------------ #
-# @spaces.GPU(durations=300)
-# def infer(edit_images, prompt, mask_prompt,
-#           seed=42, randomize_seed=False, width=1024, height=1024,
-#           guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
-#     # Get the base image from the "background" layer.
-#     image = edit_images["background"]
-#     width, height = calculate_optimal_dimensions(image)
-
-#     # If a mask prompt is provided, use the SAM-based mask generator.
-#     if mask_prompt and mask_prompt.strip() != "":
-#         try:
-#             mask = generate_mask_with_sam(image, mask_prompt)
-#         except Exception as e:
-#             raise ValueError("Error generating mask from prompt: " + str(e))
-#     else:
-#         # Fall back to using a manually drawn mask (from the first layer).
-#         try:
-#             mask = edit_images["layers"][0]
-#         except (TypeError, IndexError):
-#             raise ValueError("No mask provided. Please either draw a mask or supply a mask prompt.")
-
-#     if randomize_seed:
-#         seed = random.randint(0, MAX_SEED)
-
-#     # Run the inpainting diffusion pipeline with the provided prompt and mask.
-#     image_out = pipe(
-#         prompt=prompt,
-#         image=image,
-#         mask_image=mask,
-#         height=height,
-#         width=width,
-#         guidance_scale=guidance_scale,
-#         num_inference_steps=num_inference_steps,
-#         generator=torch.Generator(device='cuda').manual_seed(seed),
-#     ).images[0]
-
-#     output_image_jpg = image_out.convert("RGB")
-#     output_image_jpg.save("output.jpg", "JPEG")
-#     return output_image_jpg, seed
-
-# # ------------------ Gradio UI ------------------ #
-# css = """
-# #col-container {
-#     margin: 0 auto;
-#     max-width: 1000px;
-# }
-# """
-
-# with gr.Blocks(css=css) as demo:
-#     with gr.Column(elem_id="col-container"):
-#         gr.Markdown("# FLUX.1 [dev] with SAM (Transformers) Mask Generation")
-#         with gr.Row():
-#             with gr.Column():
-#                 # The image editor now allows you to optionally draw a mask.
-#                 edit_image = gr.ImageEditor(
-#                     label='Upload Image (and optionally draw a mask)',
-#                     type='pil',
-#                     sources=["upload", "webcam"],
-#                     image_mode='RGB',
-#                     layers=False,  # We will generate a mask automatically if needed.
-#                     brush=gr.Brush(colors=["#FFFFFF"]),
-#                 )
-#                 prompt = gr.Text(
-#                     label="Inpainting Prompt",
-#                     show_label=False,
-#                     max_lines=2,
-#                     placeholder="Enter your inpainting prompt",
-#                     container=False,
-#                 )
-#                 mask_prompt = gr.Text(
-#                     label="Mask Prompt (enter a coordinate as 'x,y')",
-#                     show_label=True,
-#                     placeholder="E.g. 450,600",
-#                     container=True,
-#                 )
-#                 generate_mask_btn = gr.Button("Generate Mask")
-#                 mask_preview = gr.Image(label="Mask Preview", show_label=True)
-#                 run_button = gr.Button("Run")
-#             result = gr.Image(label="Result", show_label=False)
-        
-#         # Button to preview the generated mask.
-#         def on_generate_mask(image, mask_prompt):
-#             if image is None or mask_prompt.strip() == "":
-#                 return None
-#             mask = generate_mask_with_sam(image, mask_prompt)
-#             return mask
-        
-#         generate_mask_btn.click(
-#             fn=on_generate_mask,
-#             inputs=[edit_image, mask_prompt],
-#             outputs=[mask_preview]
-#         )
-
-#         with gr.Accordion("Advanced Settings", open=False):
-#             seed = gr.Slider(
-#                 label="Seed",
-#                 minimum=0,
-#                 maximum=MAX_SEED,
-#                 step=1,
-#                 value=0,
-#             )
-#             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-#             with gr.Row():
-#                 width = gr.Slider(
-#                     label="Width",
-#                     minimum=256,
-#                     maximum=MAX_IMAGE_SIZE,
-#                     step=32,
-#                     value=1024,
-#                     visible=False
-#                 )
-#                 height = gr.Slider(
-#                     label="Height",
-#                     minimum=256,
-#                     maximum=MAX_IMAGE_SIZE,
-#                     step=32,
-#                     value=1024,
-#                     visible=False
-#                 )
-#             with gr.Row():
-#                 guidance_scale = gr.Slider(
-#                     label="Guidance Scale",
-#                     minimum=1,
-#                     maximum=30,
-#                     step=0.5,
-#                     value=3.5,
-#                 )
-#                 num_inference_steps = gr.Slider(
-#                     label="Number of Inference Steps",
-#                     minimum=1,
-#                     maximum=50,
-#                     step=1,
-#                     value=28,
-#                 )
-
-#     gr.on(
-#         triggers=[run_button.click, prompt.submit],
-#         fn=infer,
-#         inputs=[edit_image, prompt, mask_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
-#         outputs=[result, seed]
-#     )
-
-# demo.launch()
+demo.launch()