Update app.py

app.py

@@ -1,37 +1,37 @@
 import gradio as gr
 from PIL import Image, ImageDraw, ImageFilter
-import requests
-from io import BytesIO
 import torch
+from diffusers import StableDiffusionXLImg2ImgPipeline, StableDiffusionXLInpaintPipeline
 import torchvision.transforms as T
 from torchvision import models
 import numpy as np
 import cv2
 
-# AI model repo for design generation
-repo = "artificialguybr/TshirtDesignRedmond-V2"
+# Load the SDXL pipelines
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pipe_tshirt = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
+).to(device)
 
+pipe_design = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
+).to(device)
+
+# Generate T-shirt image
 def generate_cloth(color_prompt):
-    prompt = f"A plain {color_prompt} colored T-shirt hanging on a plain wall."
-    api_url = f"https://api-inference.huggingface.co/models/{repo}"
-    headers = {}
-    payload = {"inputs": prompt, "parameters": {"num_inference_steps": 30}}
-    response = requests.post(api_url, headers=headers, json=payload)
-    if response.status_code == 200:
-        return Image.open(BytesIO(response.content)).convert("RGB")
-    else:
-        raise Exception(f"Error generating cloth: {response.status_code}")
+    prompt = f"A plain {color_prompt} T-shirt hanging on a plain wall, realistic, high-quality photo."
+    image = pipe_tshirt(
+        prompt=prompt, strength=0.8, guidance_scale=7.5, num_inference_steps=30
+    ).images[0]
+    return image
 
+# Generate design image
 def generate_design(design_prompt):
     prompt = f"A bold {design_prompt} design with vibrant colors, highly detailed."
-    api_url = f"https://api-inference.huggingface.co/models/{repo}"
-    headers = {}
-    payload = {"inputs": prompt, "parameters": {"num_inference_steps": 30}}
-    response = requests.post(api_url, headers=headers, json=payload)
-    if response.status_code == 200:
-        return Image.open(BytesIO(response.content)).convert("RGBA")
-    else:
-        raise Exception(f"Error generating design: {response.status_code}")
+    image = pipe_design(
+        prompt=prompt, strength=0.8, guidance_scale=7.5, num_inference_steps=30
+    ).images[0]
+    return image.convert("RGBA")
 
 # Load pretrained DeepLabV3 model for T-shirt segmentation
 segmentation_model = models.segmentation.deeplabv3_resnet101(pretrained=True).eval()