Upload app(1).py

app(1).py

@@ -0,0 +1,187 @@
+import gradio as gr
+import numpy as np
+import spaces
+import torch
+import random
+from PIL import Image
+
+from diffusers import FluxKontextPipeline
+from diffusers.utils import load_image
+from huggingface_hub import hf_hub_download
+from aura_sr import AuraSR
+
+# --- Main Model Initialization ---
+MAX_SEED = np.iinfo(np.int32).max
+pipe = FluxKontextPipeline.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16).to("cuda")
+
+# --- Load All Adapters ---
+pipe.load_lora_weights("prithivMLmods/PhotoCleanser-i2i", weight_name="PhotoCleanser-i2i.safetensors", adapter_name="cleanser")
+pipe.load_lora_weights("prithivMLmods/Photo-Restore-i2i", weight_name="Photo-Restore-i2i.safetensors", adapter_name="restorer")
+pipe.load_lora_weights("prithivMLmods/Polaroid-Warm-i2i", weight_name="Polaroid-Warm-i2i.safetensors", adapter_name="polaroid")
+pipe.load_lora_weights("prithivMLmods/Monochrome-Pencil", weight_name="Monochrome-Pencil-i2i.safetensors", adapter_name="pencil")
+
+# --- Upscaler Model Initialization ---
+# FIX: Removed the .to("cuda") call. The library handles device placement automatically.
+aura_sr = AuraSR.from_pretrained("fal/AuraSR-v2")
+
+@spaces.GPU
+def infer(input_image, prompt, lora_adapter, upscale_image, seed=42, randomize_seed=False, guidance_scale=2.5, steps=28, progress=gr.Progress(track_tqdm=True)):
+    """
+    Perform image editing and optional upscaling.
+    
+    Args:
+        input_image (PIL.Image.Image): The input image to be edited.
+        prompt (str): Text description of the desired edit.
+        lora_adapter (str): The name of the LoRA adapter to use.
+        upscale_image (bool): If True, the final image will be upscaled 4x.
+        seed (int, optional): Random seed for reproducible generation.
+        randomize_seed (bool, optional): If True, generates a random seed.
+        guidance_scale (float, optional): Controls adherence to the prompt.
+        steps (int, optional): Number of diffusion steps.
+        progress (gr.Progress, optional): Gradio progress tracker.
+    
+    Returns:
+        tuple: A 3-tuple containing the generated/upscaled image, the seed used, 
+               and a Gradio update to make the reuse button visible.
+    """
+    if lora_adapter == "PhotoCleanser":
+        pipe.set_adapters(["cleanser"], adapter_weights=[1.0])
+    elif lora_adapter == "PhotoRestorer":
+        pipe.set_adapters(["restorer"], adapter_weights=[1.0])
+    elif lora_adapter == "PolaroidWarm":
+        pipe.set_adapters(["polaroid"], adapter_weights=[1.0])
+    elif lora_adapter == "MonochromePencil":
+        pipe.set_adapters(["pencil"], adapter_weights=[1.0])
+
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    
+    if input_image:
+        input_image = input_image.convert("RGB")
+        image = pipe(
+            image=input_image, 
+            prompt=prompt,
+            guidance_scale=guidance_scale,
+            width = input_image.size[0],
+            height = input_image.size[1],
+            num_inference_steps=steps,
+            generator=torch.Generator().manual_seed(seed),
+        ).images[0]
+    else:
+        image = pipe(
+            prompt=prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=steps,
+            generator=torch.Generator().manual_seed(seed),
+        ).images[0]
+
+    # Conditionally upscale the generated image
+    if upscale_image:
+        progress(0.8, desc="Upscaling image...")
+        image = aura_sr.upscale_4x(image)
+
+    return image, seed, gr.Button(visible=True)
+
+@spaces.GPU
+def infer_example(input_image, prompt, lora_adapter):
+    """
+    Wrapper function for gr.Examples. Upscaling is disabled for examples for speed.
+    """
+    image, seed, _ = infer(input_image, prompt, lora_adapter, upscale_image=False)
+    return image, seed
+
+css="""
+#col-container {
+    margin: 0 auto;
+    max-width: 960px;
+}
+"""
+
+with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
+    
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(f"""# **[Photo-Mate-i2i](https://huggingface.co/collections/prithivMLmods/i2i-kontext-exp-68ce573b5c0623476b636ec7)**
+        Image manipulation with Kontext adapters""")
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(label="Upload the image for editing", type="pil", height="300")
+                with gr.Row():
+                    prompt = gr.Text(
+                        label="Prompt",
+                        show_label=False,
+                        max_lines=1,
+                        placeholder="Enter your prompt for editing (e.g., 'Remove glasses', 'Add a hat')",
+                        container=False,
+                    )
+                    run_button = gr.Button("Run", scale=0)
+                with gr.Accordion("Advanced Settings", open=False):
+                    
+                    upscale_checkbox = gr.Checkbox(label="Upscale the final image", value=False)
+                    
+                    seed = gr.Slider(
+                        label="Seed",
+                        minimum=0,
+                        maximum=MAX_SEED,
+                        step=1,
+                        value=0,
+                    )
+                    
+                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+                    
+                    guidance_scale = gr.Slider(
+                        label="Guidance Scale",
+                        minimum=1,
+                        maximum=10,
+                        step=0.1,
+                        value=2.5,
+                    )       
+                    
+                    steps = gr.Slider(
+                        label="Steps",
+                        minimum=1,
+                        maximum=30,
+                        value=28,
+                        step=1
+                    )
+                    
+            with gr.Column():
+                result = gr.Image(label="Result", show_label=False, interactive=False, format="png")
+                reuse_button = gr.Button("Reuse this image", visible=False)
+                with gr.Row():
+                    lora_adapter = gr.Dropdown(
+                    label="Choose Adapter",
+                    choices=["PhotoCleanser", "PhotoRestorer", "PolaroidWarm", "MonochromePencil"],
+                    value="PhotoCleanser"
+                )
+                    
+        gr.Examples(
+            examples=[
+                ["photocleanser/1.png", "[photo content], remove the embroidered pattern from the image while preserving the background and remaining elements, maintaining realism and original details.", "PhotoCleanser"],
+                ["photocleanser/2.png", "[photo content], remove the cat from the image while preserving the background and remaining elements, maintaining realism and original details.", "PhotoCleanser"],
+                ["photorestore/1.png", "[photo content], restore and enhance the image by repairing any damage, scratches, or fading. Colorize the photo naturally while preserving authentic textures and details, maintaining a realistic and historically accurate look.", "PhotoRestorer"],
+                ["photorestore/2.png", "[photo content], restore and enhance the image by repairing any damage, scratches, or fading. Colorize the photo naturally while preserving authentic textures and details, maintaining a realistic and historically accurate look.", "PhotoRestorer"],
+                ["polaroid/1.png", "[photo content], apply a warm, vintage Polaroid-style filter, enhancing the image with nostalgic tones, soft focus, and characteristic light leaks for an authentic, retro feel.", "PolaroidWarm"],
+                ["polaroid/2.png", "[photo content], give the image a classic Polaroid look with warm, saturated colors, gentle fading, and a subtle vignette effect, evoking a sense of timeless memories.", "PolaroidWarm"],
+                ["pencil/1.png", "[photo content], transform the image into a detailed monochrome pencil sketch, emphasizing sharp lines, textures, and shading for a classic hand-drawn look.", "MonochromePencil"],
+                ["pencil/2.png", "[photo content], convert the photo into a realistic graphite pencil drawing, capturing the subject's form and depth with varied strokes and contrast.", "MonochromePencil"]
+            ],
+            inputs=[input_image, prompt, lora_adapter],
+            outputs=[result, seed],
+            fn=infer_example,
+            cache_examples=False,
+            label="Examples (Image | Prompt | Selected LoRA)"
+        )
+            
+    gr.on(
+        triggers=[run_button.click, prompt.submit],
+        fn=infer,
+        inputs=[input_image, prompt, lora_adapter, upscale_checkbox, seed, randomize_seed, guidance_scale, steps],
+        outputs=[result, seed, reuse_button]
+    )
+    reuse_button.click(
+        fn=lambda image: image,
+        inputs=[result],
+        outputs=[input_image]
+    )
+
+demo.launch(mcp_server=True)