Create app.py

app.py

@@ -0,0 +1,198 @@
+import os
+import torch
+from PIL import Image
+import numpy as np
+import traceback
+import gradio as gr
+from transformers import DetrImageProcessor, DetrForObjectDetection, LayoutLMTokenizer, LayoutLMForTokenClassification
+from diffusers import StableDiffusionPipeline, StableDiffusionUpscalePipeline
+from huggingface_hub import login
+import torchvision.transforms as T
+import torchvision.models as models
+from dotenv import load_dotenv
+
+# Load environment variables from .env file
+load_dotenv()
+
+# Retrieve Hugging Face token from environment variable
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+## 2.1 Image Analysis with DETR
+def load_detr_model():
+    try:
+        detr_model = DetrForObjectDetection.from_pretrained('facebook/detr-resnet-50')
+        detr_processor = DetrImageProcessor.from_pretrained('facebook/detr-resnet-50')
+        return detr_model, detr_processor, None
+    except Exception as e:
+        return None, None, f"Error loading DETR model: {e}"
+
+detr_model, detr_processor, detr_error = load_detr_model()
+
+def detect_objects(image):
+    if detr_model is not None and detr_processor is not None:
+        try:
+            inputs = detr_processor(images=image, return_tensors="pt")
+            outputs = detr_model(**inputs)
+            target_sizes = torch.tensor([image.size[::-1]])
+            results = detr_processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
+            return results, None
+        except Exception as e:
+            return None, f"Error in detect_objects: {e}"
+    else:
+        return None, "DETR models not loaded. Skipping object detection."
+
+## 2.2 Style Transfer with Deep Image Prior
+def style_transfer(content_image, style_image):
+    try:
+        transform = T.Compose([
+            T.Resize((512, 512)),
+            T.ToTensor(),
+            T.Lambda(lambda x: x.mul(255))
+        ])
+
+        content = transform(content_image).unsqueeze(0).requires_grad_(False)
+        style = transform(style_image).unsqueeze(0).requires_grad_(False)
+
+        vgg = models.vgg19(pretrained=True).features.eval()
+        for param in vgg.parameters():
+            param.requires_grad_(False)
+
+        generated = content.clone().requires_grad_(True)
+        optimizer = torch.optim.Adam([generated], lr=0.003)
+
+        for i in range(300):
+            generated_features = vgg(generated)
+            content_features = vgg(content)
+            style_features = vgg(style)
+
+            content_loss = torch.mean((generated_features - content_features)**2)
+            style_loss = torch.mean((generated_features - style_features)**2)
+
+            total_loss = content_loss + style_loss
+            optimizer.zero_grad()
+            total_loss.backward()
+            optimizer.step()
+
+        generated_image = generated.squeeze().clamp(0, 255).cpu().detach().numpy().transpose(1, 2, 0)
+        return Image.fromarray(np.uint8(generated_image)), None
+    except Exception as e:
+        return content_image, f"Error in style_transfer: {e}"
+
+## 2.3 Layout Generation with LayoutLM
+def load_layoutlm_model():
+    try:
+        layoutlm_tokenizer = LayoutLMTokenizer.from_pretrained('microsoft/layoutlm-base-uncased')
+        layoutlm_model = LayoutLMForTokenClassification.from_pretrained('microsoft/layoutlm-base-uncased')
+        return layoutlm_tokenizer, layoutlm_model, None
+    except Exception as e:
+        return None, None, f"Error loading LayoutLM model: {e}"
+
+layoutlm_tokenizer, layoutlm_model, layoutlm_error = load_layoutlm_model()
+
+def generate_layout(text):
+    if layoutlm_tokenizer is not None and layoutlm_model is not None:
+        try:
+            inputs = layoutlm_tokenizer(text, return_tensors="pt")
+            outputs = layoutlm_model(**inputs)
+            layout = outputs.logits.argmax(dim=-1)
+            return layout, None
+        except Exception as e:
+            return None, f"Error in generate_layout: {e}"
+    else:
+        return None, "LayoutLM models not loaded. Skipping layout generation."
+
+## 2.4 Image Generation with Stable Diffusion
+def load_stable_diffusion_model():
+    try:
+        if HF_TOKEN is None:
+            raise ValueError("Hugging Face token not found in environment variables.")
+        login(token=HF_TOKEN)
+        sd_pipeline = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4").to("cuda")
+        return sd_pipeline, None
+    except Exception as e:
+        return None, f"Error loading Stable Diffusion model: {e}"
+
+sd_pipeline, sd_error = load_stable_diffusion_model()
+
+def generate_image(prompt):
+    if sd_pipeline is not None:
+        try:
+            image = sd_pipeline(prompt).images[0]
+            return image, None
+        except Exception as e:
+            return None, f"Error in generate_image: {e}"
+    else:
+        return None, "Stable Diffusion model not loaded. Skipping image generation."
+
+## 2.5 Super-Resolution
+def load_upscale_pipeline():
+    try:
+        upscale_pipeline = StableDiffusionUpscalePipeline.from_pretrained("stabilityai/stable-diffusion-x4-upscaler").to("cuda")
+        return upscale_pipeline, None
+    except Exception as e:
+        return None, f"Error loading Upscale Pipeline: {e}"
+
+upscale_pipeline, upscale_error = load_upscale_pipeline()
+
+def super_resolve(image):
+    if upscale_pipeline is not None:
+        try:
+            if not isinstance(image, Image.Image):
+                raise ValueError("Input must be a PIL image.")
+            upscaled_image = upscale_pipeline(image=image).images[0]
+            return upscaled_image, None
+        except Exception as e:
+            return None, f"Error in super_resolve: {e}"
+    else:
+        return image, "Upscale Pipeline not loaded. Skipping super-resolution."
+
+# Step 3: Gradio Interface and Integration
+def process_image(image, style_image, text_prompt):
+    try:
+        # Detect objects
+        object_results, detect_error = detect_objects(image)
+        if detect_error:
+            return None, detect_error
+
+        # Style transfer
+        styled_image, style_error = style_transfer(image, style_image)
+        if style_error:
+            return None, style_error
+
+        # Generate layout
+        layout_results, layout_error = generate_layout(text_prompt)
+        if layout_error:
+            return None, layout_error
+
+        # Generate image based on layout
+        generated_image, gen_image_error = generate_image("modern interior design based on layout")
+        if gen_image_error:
+            return None, gen_image_error
+
+        # Super-resolve the generated image
+        final_image, upscale_error = super_resolve(generated_image)
+        if upscale_error:
+            return None, upscale_error
+
+        return final_image, None
+    except Exception as e:
+        return None, f"Error in process_image: {e}"
+
+iface = gr.Interface(
+    fn=process_image,
+    inputs=[
+        gr.Image(type="pil", label="Upload Room Image"),
+        gr.Image(type="pil", label="Upload Style Image"),
+        gr.Textbox(label="Enter Design Prompt")
+    ],
+    outputs=[
+        gr.Image(type="pil", label="Generated Image"),
+        gr.Textbox(label="Error Message")
+    ]
+)
+
+try:
+    iface.launch()
+except Exception as e:
+    print(f"Error occurred while launching the interface: {e}")
+    traceback.print_exc()