Update app.py

app.py

@@ -1,431 +1,35 @@
-import gradio as gr
-import spaces
-import torch
-import diffusers
-import transformers
-import copy
-import random
-import numpy as np
-import torchvision.transforms as T
-import math
 import os
-import peft
-from peft import LoraConfig
-from safetensors import safe_open
-from omegaconf import OmegaConf
-from omnitry.models.transformer_flux import FluxTransformer2DModel
-from omnitry.pipelines.pipeline_flux_fill import FluxFillPipeline
-from PIL import Image
-
-from huggingface_hub import snapshot_download
-snapshot_download(repo_id="Kunbyte/OmniTry", local_dir="./OmniTry")
-
-device = torch.device('cuda:0')
-weight_dtype = torch.bfloat16
-args = OmegaConf.load('configs/omnitry_v1_unified.yaml')
-
-# init model
-transformer = FluxTransformer2DModel.from_pretrained(
-    'black-forest-labs/FLUX.1-Fill-dev',
-    subfolder='transformer'
-).requires_grad_(False).to(device, dtype=weight_dtype)
-
-pipeline = FluxFillPipeline.from_pretrained(
-    'black-forest-labs/FLUX.1-Fill-dev',
-    transformer=transformer,
-    torch_dtype=weight_dtype
-).to(device)
-
-# insert LoRA
-lora_config = LoraConfig(
-    r=args.lora_rank,
-    lora_alpha=args.lora_alpha,
-    init_lora_weights="gaussian",
-    target_modules=[
-        'x_embedder',
-        'attn.to_k', 'attn.to_q', 'attn.to_v', 'attn.to_out.0',
-        'attn.add_k_proj', 'attn.add_q_proj', 'attn.add_v_proj', 'attn.to_add_out',
-        'ff.net.0.proj', 'ff.net.2', 'ff_context.net.0.proj', 'ff_context.net.2',
-        'norm1_context.linear', 'norm1.linear', 'norm.linear', 'proj_mlp', 'proj_out'
-    ]
-)
-transformer.add_adapter(lora_config, adapter_name='vtryon_lora')
-transformer.add_adapter(lora_config, adapter_name='garment_lora')
-
-with safe_open('OmniTry/omnitry_v1_unified.safetensors', framework="pt") as f:
-    lora_weights = {k: f.get_tensor(k) for k in f.keys()}
-    transformer.load_state_dict(lora_weights, strict=False)
-
-# hack lora forward
-def create_hacked_forward(module):
-
-    def lora_forward(self, active_adapter, x, *args, **kwargs):
-        result = self.base_layer(x, *args, **kwargs)
-        if active_adapter is not None:
-            lora_A = self.lora_A[active_adapter]
-            lora_B = self.lora_B[active_adapter]
-            dropout = self.lora_dropout[active_adapter]
-            scaling = self.scaling[active_adapter]
-            x = x.to(lora_A.weight.dtype)
-            result = result + lora_B(lora_A(dropout(x))) * scaling
-        return result
-
-    def hacked_lora_forward(self, x, *args, **kwargs):
-        return torch.cat((
-            lora_forward(self, 'vtryon_lora', x[:1], *args, **kwargs),
-            lora_forward(self, 'garment_lora', x[1:], *args, **kwargs),
-        ), dim=0)
-
-    return hacked_lora_forward.__get__(module, type(module))
-
-for n, m in transformer.named_modules():
-    if isinstance(m, peft.tuners.lora.layer.Linear):
-        m.forward = create_hacked_forward(m)
-
-
-def seed_everything(seed=0):
-    random.seed(seed)
-    os.environ['PYTHONHASHSEED'] = str(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-
-
-# Category mapping with sample images
-CATEGORY_SAMPLES = {
-    'top': 'top.png',
-    'bottom': 'bottom.png',
-    'dress': 'dress.jpg',
-    'hat': 'hat.png',
-    'sunglasses': 'sunglasses.png',
-    'glasses': 'glasses.png',
-    'necklace': 'necklace.png',
-    'earrings': 'earrings.png',
-    'bracelet': 'bracelet.png',
-    'ring': 'ring.png',
-    'tie': 'tie.png',
-    'bow tie': 'bow tie.png',
-    'belt': 'belt.png',
-    'shoe': 'shoe.png',
-    'bag': 'bag.png'
-}
-
-# Person sample images
-PERSON_SAMPLES = ['woman.png', 'man.png']
-
-def load_sample_image(category):
-    """Load sample image for selected category"""
-    if category and category in CATEGORY_SAMPLES:
-        img_path = CATEGORY_SAMPLES[category]
-        if os.path.exists(img_path):
-            return Image.open(img_path)
-    return None
-
-def load_random_person():
-    """Load random person image on initialization"""
-    person_img = random.choice(PERSON_SAMPLES)
-    if os.path.exists(person_img):
-        return Image.open(person_img)
-    return None
-
-def create_category_html():
-    """Create HTML for category thumbnails"""
-    html = '<div class="category-container">'
-    for category, img_file in CATEGORY_SAMPLES.items():
-        if os.path.exists(img_file):
-            # Create base64 encoded thumbnail for each category
-            import base64
-            from io import BytesIO
-            img = Image.open(img_file)
-            # Resize for thumbnail
-            img.thumbnail((80, 80), Image.Resampling.LANCZOS)
-            buffered = BytesIO()
-            img.save(buffered, format="PNG")
-            img_str = base64.b64encode(buffered.getvalue()).decode()
-            
-            html += f'''
-            <div class="category-item" data-category="{category}">
-                <img src="data:image/png;base64,{img_str}" alt="{category}">
-                <span>{category.title()}</span>
-            </div>
-            '''
-    html += '</div>'
-    return html
-
-
-@spaces.GPU
-def generate(person_image, object_image, object_class, steps, guidance_scale, seed):
-    if seed == -1:
-        seed = random.randint(0, 2**32 - 1)
-    seed_everything(seed)
-
-    max_area = 1024 * 1024
-    oW, oH = person_image.width, person_image.height
-    ratio = math.sqrt(max_area / (oW * oH))
-    ratio = min(1, ratio)
-    tW, tH = int(oW * ratio) // 16 * 16, int(oH * ratio) // 16 * 16
-    transform = T.Compose([
-        T.Resize((tH, tW)),
-        T.ToTensor(),
-    ])
-    person_image = transform(person_image)
-
-    ratio = min(tW / object_image.width, tH / object_image.height)
-    transform = T.Compose([
-        T.Resize((int(object_image.height * ratio), int(object_image.width * ratio))),
-        T.ToTensor(),
-    ])
-    object_image_padded = torch.ones_like(person_image)
-    object_image = transform(object_image)
-    new_h, new_w = object_image.shape[1], object_image.shape[2]
-    min_x = (tW - new_w) // 2
-    min_y = (tH - new_h) // 2
-    object_image_padded[:, min_y: min_y + new_h, min_x: min_x + new_w] = object_image
-
-    prompts = [args.object_map[object_class]] * 2
-    img_cond = torch.stack([person_image, object_image_padded]).to(dtype=weight_dtype, device=device)
-    mask = torch.zeros_like(img_cond).to(img_cond)
-
-    with torch.no_grad():
-        img = pipeline(
-            prompt=prompts,
-            height=tH,
-            width=tW,
-            img_cond=img_cond,
-            mask=mask,
-            guidance_scale=guidance_scale,
-            num_inference_steps=steps,
-            generator=torch.Generator(device).manual_seed(seed),
-        ).images[0]
-
-    return img
-
-
-# Custom CSS
-custom_css = """
-/* 전체 배경 */
-.gradio-container {
-    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
-    font-family: 'Inter', sans-serif;
-}
-
-/* Category thumbnails container */
-.category-container {
-    display: flex;
-    flex-wrap: wrap;
-    gap: 15px;
-    justify-content: center;
-    margin: 20px 0;
-    padding: 20px;
-    background: rgba(255, 255, 255, 0.1);
-    border-radius: 15px;
-    backdrop-filter: blur(10px);
-}
-
-.category-item {
-    display: flex;
-    flex-direction: column;
-    align-items: center;
-    cursor: pointer;
-    padding: 10px;
-    border-radius: 10px;
-    background: rgba(255, 255, 255, 0.9);
-    transition: all 0.3s ease;
-    min-width: 90px;
-}
-
-.category-item:hover {
-    transform: translateY(-5px);
-    box-shadow: 0 5px 20px rgba(0, 0, 0, 0.2);
-    background: white;
-}
-
-.category-item.selected {
-    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
-    color: white;
-}
-
-.category-item img {
-    width: 60px;
-    height: 60px;
-    object-fit: contain;
-    margin-bottom: 5px;
-}
-
-.category-item span {
-    font-size: 0.85em;
-    text-align: center;
-    font-weight: 500;
-}
-
-/* === 플레이스홀더 전부 제거 === */
-.gr-image svg,
-.gr-image [data-testid*="placeholder"],
-.gr-image [class*="placeholder"],
-.gr-image [aria-label*="placeholder"],
-.gr-image [class*="svelte"][class*="placeholder"],
-.gr-image .absolute.inset-0.flex.items-center.justify-center,
-.gr-image .flex.items-center.justify-center svg {
-    display: none !important;
-    visibility: hidden !important;
-}
-.gr-image [class*="overlay"],
-.gr-image .fixed.inset-0,
-.gr-image .absolute.inset-0 {
-    pointer-events: none !important;
-}
-
-/* 이미지 업로드 영역 */
-.gr-image .wrap { background: transparent !important; min-height: 400px !important; }
-.gr-image .upload-container {
-    min-height: 400px !important;
-    border: 3px dashed rgba(102, 126, 234, 0.4) !important;
-    border-radius: 12px !important;
-    background: linear-gradient(135deg, rgba(248, 250, 252, 0.5) 0%, rgba(241, 245, 249, 0.5) 100%) !important;
-    position: relative !important;
-}
-/* 이미지 있을 때 */
-.gr-image:has(img) .upload-container { border: none !important; background: transparent !important; }
-/* 안내 텍스트 */
-.gr-image .upload-container::after {
-    content: "Click or Drag to Upload";
-    position: absolute; top: 50%; left: 50%;
-    transform: translate(-50%, -50%);
-    color: rgba(102, 126, 234, 0.7);
-    font-size: 1.05em; font-weight: 500;
-    pointer-events: none;
-}
-.gr-image:has(img) .upload-container::after { display: none !important; }
-/* 업로드 이미지 */
-.gr-image img { border-radius: 12px !important; position: relative !important; z-index: 10 !important; }
-
-/* 버튼 스타일 */
-.gr-button-primary {
-    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
-    color: white !important; border: none !important;
-    padding: 15px 40px !important; font-size: 1.2em !important;
-    border-radius: 50px !important; cursor: pointer !important;
-}
-
-/* Radio button styling for categories */
-.gr-radio {
-    display: none !important;
-}
-"""
-
-if __name__ == '__main__':
-
-    with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
-        with gr.Column(elem_id="header"):
-            gr.HTML("""
-                <h1 style="text-align: center; color: white; margin: 20px 0;">
-                    ✨ CodiFit-AI Virtual Try-On ✨
-                </h1>
-                <p style="text-align: center; color: rgba(255,255,255,0.9); margin-bottom: 30px;">
-                    Experience the future of fashion with AI-powered virtual clothing try-on
-                </p>
-            """)
-            
-            # Category selector with thumbnails
-            gr.Markdown("### Select Fashion Category")
-            with gr.Row():
-                category_selector = gr.Radio(
-                    choices=list(CATEGORY_SAMPLES.keys()),
-                    value='top',
-                    label="Fashion Categories",
-                    elem_id="category_radio",
-                    visible=False  # Hide default radio buttons
-                )
-            
-            # Display category thumbnails
-            category_html = gr.HTML(create_category_html())
-            
-        with gr.Row(equal_height=True):
-            with gr.Column(scale=1):
-                person_image = gr.Image(
-                    type="pil", 
-                    label="Upload Person Photo", 
-                    height=500, 
-                    interactive=True,
-                    value=load_random_person()  # Load random person on init
-                )
-                
-            with gr.Column(scale=1):
-                object_image = gr.Image(
-                    type="pil", 
-                    label="Upload Object Image", 
-                    height=400, 
-                    interactive=True,
-                    value=load_sample_image('top')  # Load default top image
-                )
-                object_class = gr.Dropdown(
-                    label='Selected Object Category', 
-                    choices=list(args.object_map.keys()),
-                    value='top',
-                    interactive=True
-                )
-                run_button = gr.Button(value="🚀 Generate Try-On", variant='primary')
-                
-            with gr.Column(scale=1):
-                image_out = gr.Image(type="pil", label="Virtual Try-On Result", height=500, interactive=False)
-
-        with gr.Accordion("⚙️ Advanced Settings", open=False):
-            with gr.Row():
-                guidance_scale = gr.Slider(label="🎯 Guidance Scale", minimum=1, maximum=50, value=30, step=0.1)
-                steps = gr.Slider(label="🔄 Inference Steps", minimum=1, maximum=50, value=20, step=1)
-                seed = gr.Number(label="🎲 Random Seed", value=-1, precision=0)
-
-        # JavaScript for category selection interaction
-        demo.load(None, None, None, js="""
-        function() {
-            // Add click handlers to category items
-            document.querySelectorAll('.category-item').forEach(item => {
-                item.addEventListener('click', function() {
-                    // Remove selected class from all items
-                    document.querySelectorAll('.category-item').forEach(i => 
-                        i.classList.remove('selected')
-                    );
-                    // Add selected class to clicked item
-                    this.classList.add('selected');
-                    
-                    // Get category name
-                    const category = this.dataset.category;
-                    
-                    // Trigger category selection (would need proper Gradio event handling)
-                    console.log('Selected category:', category);
-                });
-            });
+import sys
+import streamlit as st
+from tempfile import NamedTemporaryFile
+
+def main():
+    try:
+        # Get the code from secrets
+        code = os.environ.get("MAIN_CODE")
+        
+        if not code:
+            st.error("⚠️ The application code wasn't found in secrets. Please add the MAIN_CODE secret.")
+            return
+        
+        # Create a temporary Python file
+        with NamedTemporaryFile(suffix='.py', delete=False, mode='w') as tmp:
+            tmp.write(code)
+            tmp_path = tmp.name
+        
+        # Execute the code
+        exec(compile(code, tmp_path, 'exec'), globals())
+        
+        # Clean up the temporary file
+        try:
+            os.unlink(tmp_path)
+        except:
+            pass
             
-            // Select first category by default
-            document.querySelector('.category-item')?.classList.add('selected');
-        }
-        """)
-
-        # Handle category selection
-        def on_category_select(category):
-            """Update object image and class when category is selected"""
-            sample_img = load_sample_image(category)
-            return sample_img, category
-
-        # Connect category selector to update object image and class
-        category_selector.change(
-            fn=on_category_select,
-            inputs=[category_selector],
-            outputs=[object_image, object_class]
-        )
-
-        # Manual category selection through radio (hidden but functional)
-        def update_from_thumbnail(evt: gr.SelectData):
-            category = list(CATEGORY_SAMPLES.keys())[evt.index]
-            sample_img = load_sample_image(category)
-            return sample_img, category
-
-        # Run generation
-        run_button.click(
-            generate,
-            inputs=[person_image, object_image, object_class, steps, guidance_scale, seed],
-            outputs=[image_out]
-        )
+    except Exception as e:
+        st.error(f"⚠️ Error loading or executing the application: {str(e)}")
+        import traceback
+        st.code(traceback.format_exc())
 
-    demo.launch(server_name="0.0.0.0")
+if __name__ == "__main__":
+    main()