Create app.py

app.py

@@ -0,0 +1,128 @@
+import gradio as gr
+import torch
+import spaces
+from diffusers import FluxInpaintPipeline
+from PIL import Image
+
+# Initialize the pipeline
+pipe = FluxInpaintPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-dev", 
+    torch_dtype=torch.bfloat16
+)
+pipe.to("cuda")
+pipe.load_lora_weights(
+    "ali-vilab/In-Context-LoRA", 
+    weight_name="visual-identity-design.safetensors"
+)
+
+def square_center_crop(img, target_size=768):
+    if img.mode in ('RGBA', 'P'):
+        img = img.convert('RGB')
+    
+    width, height = img.size
+    crop_size = min(width, height)
+    
+    left = (width - crop_size) // 2
+    top = (height - crop_size) // 2
+    right = left + crop_size
+    bottom = top + crop_size
+    
+    img_cropped = img.crop((left, top, right, bottom))
+    return img_cropped.resize((target_size, target_size), Image.Resampling.LANCZOS)
+
+def duplicate_horizontally(img):
+    width, height = img.size
+    if width != height:
+        raise ValueError(f"Input image must be square, got {width}x{height}")
+    
+    new_image = Image.new('RGB', (width * 2, height))
+    new_image.paste(img, (0, 0))
+    new_image.paste(img, (width, 0))
+    return new_image
+
+# Load the mask image
+mask = Image.open("mask_square.png")
+
+@spaces.GPU
+def generate(image, prompt_user):
+    prompt_structure = "The two-panel image showcases the logo of a brand, [LEFT] the left panel is showing the logo [RIGHT] the right panel has this logo applied to "
+    prompt = prompt_structure + prompt_user
+    
+    cropped_image = square_center_crop(image)
+    logo_dupli = duplicate_horizontally(cropped_image)
+    
+    out = pipe(
+        prompt=prompt,
+        image=logo_dupli,
+        mask_image=mask,
+        guidance_scale=6,
+        height=768,
+        width=1536,
+        num_inference_steps=28,
+        max_sequence_length=256,
+        strength=1
+    ).images[0]
+    
+    width, height = out.size
+    half_width = width // 2
+    image_2 = out.crop((half_width, 0, width, height))
+    return image_2
+
+def process_image(input_image, prompt):
+    try:
+        if input_image is None:
+            return None, "Please upload an image first."
+        
+        if not prompt:
+            return None, "Please provide a prompt."
+            
+        result = generate(input_image, prompt)
+        return result, "Generation completed successfully!"
+    except Exception as e:
+        return None, f"Error during generation: {str(e)}"
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Logo in Context")
+    gr.Markdown("### In-Context LoRA + Image-to-Image, apply your logo to anything")
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(
+                label="Upload Logo Image",
+                type="pil",
+                height=384
+            )
+            prompt_input = gr.Textbox(
+                label="Where should the logo be applied?",
+                placeholder="e.g., a coffee cup on a wooden table",
+                lines=2
+            )
+            generate_btn = gr.Button("Generate Application", variant="primary")
+        
+        with gr.Column():
+            output_image = gr.Image(label="Generated Application")
+            status_text = gr.Textbox(
+                label="Status",
+                interactive=False
+            )
+    
+    with gr.Row():
+        gr.Markdown("""
+        ### Instructions:
+        1. Upload a logo image (preferably square)
+        2. Describe where you'd like to see the logo applied
+        3. Click 'Generate Application' and wait for the result
+        
+        Note: The generation process might take a few moments.
+        """)
+    
+    # Set up the click event
+    generate_btn.click(
+        fn=process_image,
+        inputs=[input_image, prompt_input],
+        outputs=[output_image]
+    )
+
+# Launch the interface
+if __name__ == "__main__":
+    demo.launch()