Upload 2 files

app.py

@@ -0,0 +1,138 @@
+import streamlit as st
+import torch
+from torch import nn
+from diffusers import DDPMScheduler, UNet2DModel
+import matplotlib.pyplot as plt
+from tqdm.auto import tqdm
+
+# Reuse your existing model code
+class ClassConditionedUnet(nn.Module):
+    def __init__(self, num_classes=3, class_emb_size=12):
+        super().__init__()
+        self.class_emb = nn.Embedding(num_classes, class_emb_size)
+        self.model = UNet2DModel(
+            sample_size=64,
+            in_channels=3 + class_emb_size,
+            out_channels=3,
+            layers_per_block=2,
+            block_out_channels=(64, 128, 256, 512),
+            down_block_types=(
+                "DownBlock2D",
+                "DownBlock2D",
+                "AttnDownBlock2D",
+                "AttnDownBlock2D",
+            ),
+            up_block_types=(
+                "AttnUpBlock2D",
+                "AttnUpBlock2D",
+                "UpBlock2D",
+                "UpBlock2D",
+            ),
+        )
+
+    def forward(self, x, t, class_labels):
+        bs, ch, w, h = x.shape
+        class_cond = self.class_emb(class_labels)
+        class_cond = class_cond.view(bs, class_cond.shape[1], 1, 1).expand(bs, class_cond.shape[1], w, h)
+        net_input = torch.cat((x, class_cond), 1)
+        return self.model(net_input, t).sample
+
+@st.cache_resource
+def load_model(model_path):
+    """Load the model with caching to avoid reloading"""
+    device = 'cpu'  # For deployment, we'll use CPU
+    net = ClassConditionedUnet().to(device)
+    noise_scheduler = DDPMScheduler(num_train_timesteps=1000, beta_schedule='squaredcos_cap_v2')
+    checkpoint = torch.load(model_path, map_location='cpu')
+    net.load_state_dict(checkpoint['model_state_dict'])
+    return net, noise_scheduler
+
+def generate_mixed_faces(net, noise_scheduler, mix_weights, num_images=1):
+    """Generate faces with mixed ethnic features"""
+    device = next(net.parameters()).device
+    net.eval()
+    with torch.no_grad():
+        x = torch.randn(num_images, 3, 64, 64).to(device)
+        
+        # Get embeddings for all classes
+        emb_asian = net.class_emb(torch.zeros(num_images).long().to(device))
+        emb_indian = net.class_emb(torch.ones(num_images).long().to(device))
+        emb_european = net.class_emb(torch.full((num_images,), 2).to(device))
+        
+        progress_bar = st.progress(0)
+        for idx, t in enumerate(noise_scheduler.timesteps):
+            # Update progress bar
+            progress_bar.progress(idx / len(noise_scheduler.timesteps))
+            
+            # Mix embeddings according to weights
+            mixed_emb = (
+                mix_weights[0] * emb_asian +
+                mix_weights[1] * emb_indian +
+                mix_weights[2] * emb_european
+            )
+            
+            # Override embedding layer temporarily
+            original_forward = net.class_emb.forward
+            net.class_emb.forward = lambda _: mixed_emb
+            
+            residual = net(x, t, torch.zeros(num_images).long().to(device))
+            x = noise_scheduler.step(residual, t, x).prev_sample
+            
+            # Restore original embedding layer
+            net.class_emb.forward = original_forward
+        
+        progress_bar.progress(1.0)
+    
+    x = (x.clamp(-1, 1) + 1) / 2
+    return x
+
+def main():
+    st.title("AI Face Generator with Ethnic Features Mixing")
+    
+    # Load model
+    try:
+        net, noise_scheduler = load_model('final_model/final_diffusion_model.pt')
+    except Exception as e:
+        st.error(f"Error loading model: {str(e)}")
+        return
+    
+    # Create sliders for ethnicity percentages
+    st.subheader("Adjust Ethnicity Mix")
+    col1, col2, col3 = st.columns(3)
+    
+    with col1:
+        asian_pct = st.slider("Asian Features %", 0, 100, 33, 1)
+    with col2:
+        indian_pct = st.slider("Indian Features %", 0, 100, 33, 1)
+    with col3:
+        european_pct = st.slider("European Features %", 0, 100, 34, 1)
+    
+    # Calculate total and normalize if needed
+    total = asian_pct + indian_pct + european_pct
+    if total == 0:
+        st.warning("Total percentage cannot be 0%. Please adjust the sliders.")
+        return
+    
+    # Normalize weights to sum to 1
+    weights = [asian_pct/total, indian_pct/total, european_pct/total]
+    
+    # Display current mix
+    st.write("Current mix (normalized):")
+    st.write(f"Asian: {weights[0]:.2%}, Indian: {weights[1]:.2%}, European: {weights[2]:.2%}")
+    
+    # Generate button
+    if st.button("Generate Face"):
+        try:
+            with st.spinner("Generating face..."):
+                # Generate the image
+                generated_images = generate_mixed_faces(net, noise_scheduler, weights)
+                
+                # Convert to numpy and display
+                img = generated_images[0].permute(1, 2, 0).cpu().numpy()
+                st.image(img, caption="Generated Face", use_column_width=True)
+                
+        except Exception as e:
+            st.error(f"Error generating image: {str(e)}")
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,139 @@
+accelerate==0.34.2
+aiohappyeyeballs==2.4.0
+aiohttp==3.10.5
+aiosignal==1.3.1
+altair==5.4.1
+annotated-types==0.7.0
+anyio==4.4.0
+asttokens==2.4.1
+async-timeout==4.0.3
+attrs==24.2.0
+blinker==1.8.2
+boto3==1.35.54
+botocore==1.35.54
+cachetools==5.5.0
+certifi==2024.8.30
+charset-normalizer==3.3.2
+click==8.1.7
+colorama==0.4.6
+comm==0.2.2
+contourpy==1.3.1
+cycler==0.12.1
+databricks-api==0.9.0
+databricks-cli==0.18.0
+dataclasses==0.6
+debugpy==1.8.5
+decorator==5.1.1
+diffusers==0.31.0
+exceptiongroup==1.2.2
+executing==2.1.0
+faiss-cpu==1.8.0.post1
+filelock==3.16.0
+fonttools==4.55.3
+frozenlist==1.4.1
+fsspec==2024.9.0
+gitdb==4.0.11
+GitPython==3.1.43
+greenlet==3.1.0
+h11==0.14.0
+httpcore==1.0.5
+httpx==0.27.2
+huggingface==0.0.1
+huggingface-hub==0.24.7
+idna==3.10
+importlib_metadata==8.5.0
+ipykernel==6.29.5
+ipython==8.27.0
+jedi==0.19.1
+Jinja2==3.1.4
+jmespath==1.0.1
+joblib==1.4.2
+johnsnowlabs==5.5.0
+jsonpatch==1.33
+jsonpointer==3.0.0
+jsonschema==4.23.0
+jsonschema-specifications==2023.12.1
+jupyter_client==8.6.3
+jupyter_core==5.7.2
+kiwisolver==1.4.7
+langchain==0.3.0
+langchain-core==0.3.0
+langchain-text-splitters==0.3.0
+langsmith==0.1.121
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+matplotlib==3.9.3
+matplotlib-inline==0.1.7
+mdurl==0.1.2
+mpmath==1.3.0
+multidict==6.1.0
+narwhals==1.8.1
+nest-asyncio==1.6.0
+networkx==3.3
+nlu==5.4.1
+numpy==1.26.4
+oauthlib==3.2.2
+orjson==3.10.7
+packaging==24.1
+pandas==2.2.2
+parso==0.8.4
+pillow==10.4.0
+platformdirs==4.3.3
+prompt_toolkit==3.0.47
+protobuf==5.28.1
+psutil==6.0.0
+pure_eval==0.2.3
+py4j==0.10.9
+pyarrow==17.0.0
+pydantic==2.9.1
+pydantic_core==2.23.3
+pydeck==0.9.1
+Pygments==2.18.0
+PyJWT==2.9.0
+pyparsing==3.2.0
+pyspark==3.0.2
+python-dateutil==2.9.0.post0
+pytz==2024.2
+pywin32==306
+PyYAML==6.0.2
+pyzmq==26.2.0
+referencing==0.35.1
+regex==2024.9.11
+requests==2.32.3
+rich==13.8.1
+rpds-py==0.20.0
+s3transfer==0.10.3
+safetensors==0.4.5
+scikit-learn==1.5.2
+scipy==1.14.1
+sentence-transformers==3.1.0
+six==1.16.0
+smmap==5.0.1
+sniffio==1.3.1
+spark-nlp==5.5.0
+spark-nlp-display==5.0
+SQLAlchemy==2.0.35
+stack-data==0.6.3
+streamlit==1.38.0
+streamlit-chat==0.1.1
+svgwrite==1.4
+sympy==1.13.1
+tabulate==0.9.0
+tenacity==8.5.0
+threadpoolctl==3.5.0
+tiktoken==0.7.0
+tokenizers==0.19.1
+toml==0.10.2
+torch==2.5.1
+torchvision==0.20.1
+tornado==6.4.1
+tqdm==4.66.5
+traitlets==5.14.3
+transformers==4.44.2
+typing_extensions==4.12.2
+tzdata==2024.1
+urllib3==2.2.3
+watchdog==4.0.2
+wcwidth==0.2.13
+yarl==1.11.1
+zipp==3.21.0