Update app.py

app.py

@@ -1,59 +1,32 @@
 import gradio as gr
-import requests
-import os
+from diffusers import DiffusionPipeline
+import torch
 from PIL import Image
-from io import BytesIO
-from tqdm import tqdm
-import time
+import tempfile
 
-repo = "artificialguybr/TshirtDesignRedmond-V2"
+print("Loading Rookus T1 model...")
+pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16)
+pipe.to("cuda" if torch.cuda.is_available() else "cpu")
+pipe.set_progress_bar_config(disable=True)
 
 def infer(color_prompt, dress_type_prompt, design_prompt):
-    # Improved prompt for higher accuracy
+    # Construct prompt
     prompt = (
         f"A high-quality digital image of a {color_prompt} {dress_type_prompt}, "
-        f"featuring a {design_prompt} printed in sharp detail printedon the {dress_type_prompt},"
-        f"facing front, hanging on he plain wall"
-        f"The fabric has realistic texture,"
-        f"smooth folds, and accurate lighting. The design is perfectly aligned, with natural shadows "
-        f"and highlights, creating a photorealistic look."
+        f"featuring a {design_prompt} printed in sharp detail on the {dress_type_prompt}, "
+        f"facing front, hanging on a plain wall. "
+        f"The fabric has realistic texture, smooth folds, and accurate lighting. "
+        f"The design is perfectly aligned, with natural shadows and highlights, "
+        f"creating a photorealistic look."
     )
-    
+
     print("Generating image with prompt:", prompt)
-    api_url = f"https://api-inference.huggingface.co/models/{repo}"
-    
-    headers = {}  # If API token needed, add here
-    
-    payload = {
-        "inputs": prompt,
-        "parameters": {
-            # Optimized negative prompt
-            "negative_prompt": "low quality, artifacts, distorted, blurry, overexposed, underexposed, unrealistic texture, poor lighting, misaligned print, plastic-like fabric, grainy, washed-out colors, 3D render, cartoon, digital art, watermark, bad anatomy, malformed, cluttered design",
-            "num_inference_steps": 30,
-            "scheduler": "EulerAncestralDiscreteScheduler"  # Faster & more accurate scheduler
-        },
-    }
 
-    error_count = 0
-    pbar = tqdm(total=None, desc="Loading model")
-    
-    while True:
-        print("Sending request to API...")
-        response = requests.post(api_url, headers=headers, json=payload)
-        print("API response status code:", response.status_code)
-        
-        if response.status_code == 200:
-            print("Image generation successful!")
-            return Image.open(BytesIO(response.content))
-        elif response.status_code == 503:
-            time.sleep(1)
-            pbar.update(1)
-        elif response.status_code == 500 and error_count < 5:
-            time.sleep(1)
-            error_count += 1
-        else:
-            print("API Error:", response.status_code)
-            raise Exception(f"API Error: {response.status_code}")
+    # Generate image
+    result = pipe(prompt, num_inference_steps=4, guidance_scale=0.0)  # SDXL-Turbo is optimized for fast steps
+    image = result.images[0]
+
+    return image
 
 # Gradio Interface
 iface = gr.Interface(
@@ -70,4 +43,4 @@ iface = gr.Interface(
 )
 
 print("Launching Gradio interface...")
-iface.launch()
+iface.launch()