Update app.py

app.py

@@ -1,113 +1,112 @@
 import os
-from huggingface_hub import login
+from huggingface_hub import login, snapshot_download
 from transformers import BlipProcessor, BlipForConditionalGeneration
-from transformers import MllamaForConditionalGeneration, AutoProcessor
 from PIL import Image
 from dotenv import load_dotenv
-
 import gradio as gr
-from diffusers import DiffusionPipeline
+from diffusers import FluxPipeline
 import torch
 import spaces  # Hugging Face Spaces module
 
-import requests
-from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor
-from qwen_vl_utils import process_vision_info
+# -----------------------
+# Pre-cache models at startup
+# -----------------------
+snapshot_download("Salesforce/blip-image-captioning-large", timeout=120)
+snapshot_download("noamrot/FuseCap", timeout=120)
+snapshot_download("black-forest-labs/FLUX.1-dev", timeout=300)
+
+# -----------------------
+# Authentication
+# -----------------------
+load_dotenv()
+HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
+if HUGGINGFACE_TOKEN:
+    login(token=HUGGINGFACE_TOKEN)
 
-from diffusers import DiffusionPipeline
+# -----------------------
+# Load models
+# -----------------------
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
+processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large", timeout=120)
+model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large", timeout=120).to(device)
 
+processor1 = BlipProcessor.from_pretrained("noamrot/FuseCap", timeout=120)
+model2 = BlipForConditionalGeneration.from_pretrained("noamrot/FuseCap", timeout=120).to(device)
 
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, timeout=300).to(device)
+
+# -----------------------
+# Options
+# -----------------------
 fabrics = ['cotton', 'silk', 'denim', 'linen', 'polyester', 'wool', 'velvet']
 patterns = ['striped', 'floral', 'geometric', 'abstract', 'solid', 'polka dots']
 textile_designs = ['woven texture', 'embroidery', 'printed fabric', 'hand-dyed', 'quilting']
 
-
-
-
-# Get Hugging Face Token from environment variable
-HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
-
-# Authenticate using the token
-login(token =HUGGINGFACE_TOKEN)
-
-
-
-
-processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
-model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large")
-processor1 = BlipProcessor.from_pretrained("noamrot/FuseCap")
-model2 = BlipForConditionalGeneration.from_pretrained("noamrot/FuseCap")
-
-from diffusers import FluxPipeline
-
-pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
-
-
-
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-# pipe.to(device)
-
-model.to(device)
-pipe.to(device)
-model2.to(device)
-
-
-
+# -----------------------
+# Inference Function
+# -----------------------
 @spaces.GPU(duration=150)
 def generate_caption_and_image(image, f, p, d):
-    if f!=None and p!=None and d!=None and image!=None:
+    if image and f and p and d:
         img = image.convert("RGB")
-        # reader = easyocr.Reader(['en'])  
-        # # result = reader.readtext(img)
-        # import random
-        
-    
-    
-    
-    
-            
-    
-     
-        text = "a picture of "
-        inputs = processor(img, text, return_tensors="pt").to(device)
-        
-        out = model2.generate(**inputs, num_beams = 3)
-       
-        
-    
+
+        # Caption with FuseCap
+        inputs = processor(img, "a picture of ", return_tensors="pt").to(device)
+        out = model2.generate(**inputs, num_beams=3)
         caption2 = processor1.decode(out[0], skip_special_tokens=True)
-        
-     
+
+        # Caption with BLIP
         inputs = processor(image, return_tensors="pt", padding=True, truncation=True, max_length=250)
-        inputs = {key: val.to(device) for key, val in inputs.items()}
+        inputs = {k: v.to(device) for k, v in inputs.items()}
         out = model.generate(**inputs)
         caption1 = processor.decode(out[0], skip_special_tokens=True)
-        prompt = f"Design a high-quality, stylish clothing item that flawlessly combines the essence of {caption1} and {caption2}. The design should emphasize the luxurious feel and practicality of {f} fabric, while integrating intricate {d} textual design elements. Incorporate {p} patterns that elevate the garment's aesthetic, ensuring a harmonious blend of textures and visuals. The final piece should be both sophisticated and innovative, reflecting modern trends while preserving timeless elegance. The design should be bold, wearable, and a true work of art."
-
-    
-
-
-
 
-
-        
-    
-    
- 
-        image = pipe(prompt,height=1024,width=1024,guidance_scale=3.5,num_inference_steps=50,max_sequence_length=512,generator=torch.Generator("cpu").manual_seed(0)).images[0]
-        return image
+        # Compose prompt
+        prompt = (
+            f"Design a high-quality, stylish clothing item that combines the essence of {caption1} and {caption2}. "
+            f"Use luxurious {f} fabric with intricate {d} design elements. "
+            f"Incorporate {p} patterns to elevate the garment's aesthetic. "
+            "Ensure sophistication, innovation, and timeless elegance."
+        )
+
+        # Generate image
+        result = pipe(
+            prompt,
+            height=1024,
+            width=1024,
+            guidance_scale=3.5,
+            num_inference_steps=50,
+            max_sequence_length=512,
+            generator=torch.Generator('cpu').manual_seed(0)
+        ).images[0]
+
+        return result
     return None
+
+# -----------------------
 # Gradio UI
+# -----------------------
 iface = gr.Interface(
     fn=generate_caption_and_image,
-    inputs=[gr.Image(type="pil", label="Upload Image"), gr.Radio(fabrics, label="Select Fabric"), gr.Radio(patterns, label="Select Pattern"), gr.Radio(textile_designs, label="Select Textile Design")],
-            
-    outputs=[gr.Image(label="Generated Design 1")],
+    inputs=[
+        gr.Image(type="pil", label="Upload Image"),
+        gr.Radio(fabrics, label="Select Fabric"),
+        gr.Radio(patterns, label="Select Pattern"),
+        gr.Radio(textile_designs, label="Select Textile Design")
+    ],
+    outputs=gr.Image(label="Generated Design"),
     live=True
 )
-iface.launch(share=True)
+
+iface.launch()
+
+    
+    
+            
+
+
+